import sqlite3
import numpy as np
import glob
import re
import commands
import os


def InsertSetllarMass(halo,database="halos.db",mass_file="PopII_halos_stellarmass_0034.txt"):
    conn = sqlite3.connect(database)
    c = conn.cursor()
    c.execute("alter table Halos add column StellarMass real")    

    for line in open(mass_file):
        temp_s = line.split()
        if len(temp_s) != 3 : break  
        line = 'update Halos SET StellarMass=%f WHERE SnapHaloID=%d and SnapCurrentTimeIdentifier;'% (float(temp_s[1]),int(temp_s[0]),TimeIdentifier) 
        c.execute(line)

    conn.commit()
    c.close()

if __name__ == "__main__":
 conn = sqlite3.connect('halos.db')
 #conn.execute("CREATE INDEX halo_index ON Halos (SnapHaloID,SnapCurrentTimeIdentifier);")
 c = conn.cursor()
 #conn.execute("alter table Halos add column StellarMass real")
 #conn.execute("alter table Halos add column ActiveStarMass real")

 c.execute("PRAGMA synchronous = OFF")
 c.execute("PRAGMA journal_mode = OFF")

 os.chdir("..")
 folderlist = glob.glob('RD????')
 upperfolder = os.getcwd()
 for folder in folderlist:
   n = int(re.sub(r'\D',"",folder))
   #if n <= 5: continue
   os.chdir(upperfolder+"/"+folder)
   name,= glob.glob('RedshiftOutput????')
   a=commands.getoutput('grep CurrentTimeIdentifier '+name)
   TimeIdentifier=int(a.split()[2])
   #a=commands.getoutput('grep InitialTime '+name)
   #currenttime=float(a.split()[2])

   mass_file = "PopII_halos_stellarmass_%04i.txt"%n
   print mass_file
   conn.execute('BEGIN TRANSACTION')
   for line in open(mass_file):
        temp_s = line.split()
        if len(temp_s) != 3 : break
        li = 'update Halos SET StellarMass=%f WHERE SnapHaloID=%d and SnapCurrentTimeIdentifier=%d;'% (float(temp_s[1]),int(temp_s[0]),TimeIdentifier)
        c.execute(li)
        li = 'update Halos SET ActiveStarMass=%f WHERE SnapHaloID=%d and SnapCurrentTimeIdentifier=%d;'% (float(temp_s[2]),int(temp_s[0]),TimeIdentifier)
        c.execute(li)
   conn.commit()
   os.chdir(upperfolder)

 conn.commit()
 conn.close()