import tempfile
import swmm_ea_temp_data
import os
from  distutils.dir_util import copy_tree
import sys
import StringIO
import swmm5ec
import guiqwt
import swmmout
import swmm5.swmm5 as sw
import swmm_ea_controller
from PyQt4 import QtCore



class parameters_class_(object):
    def __init__(self):
        pass  
    def setvalues(self,attrs):
        for a in attrs:
            setattr(self, a,attrs[a])
    def __str__(self):
        s=""
        for a,b in self.__dict__.iteritems():
            s=s+ a + "\t=>\t"+str(b )+"\n"
        return s


class Project():
    
    paramfilename="param.yaml"
    LOAD_NOSWMMFILE=2
    LOAD_COMPLETE=1
    outdir="output"
    tmpdir="tmp"    

    def __init__(self, dirname=None, swmmfilename=None):
        self.setSwmmfile(None)
        self.ids=None
        self.times=None
        if dirname:
            self.dirname=dirname
            self.load(dirname,swmmfilename)
        else:
            self.dirname=tempfile.mkdtemp(prefix="swmm-ea-project", suffix="dir")
            if swmmfilename:
                self.setSwmmfile(swmmfilename)
            else:
                self.setSwmmfile("swmm.inp")
            data=swmm_ea_temp_data.data()
            paramfile=data.paramfile
            self.string_to_param(paramfile)
            self.swmm_data=data.swmmfile
            self.setSwmmfile(data.swmmfilename)
            self.save()   
        self.slotted_swmmfilename=None


    def pause_optimization(self,state):
        self.swmm5ec.pause(state)

    def stop_optimization(self):
        self.swmm5ec.stop()
        

    def run_optimization(self):
        self.swmm5ec.start()
    
    def get_ids(self):
        if not self.ids:
            self.get_swmm_ids_times()
        return self.ids
        
    def setSwmmfile(self, swmmfile=None):
        self.swmmfilename=swmmfile
        self.ids=[] # reset the ids now next time the swmm file has to be run and ides extracted. 
    
         
    def get_swmm_ids_times(self):
        # first run swmm
        sp=self.dirname+os.sep+self.swmmfilename
        binfile=sp[:-3]+"bin"
        rptfile=sp[:-3]+"rpt"  
        print "Attempting to run swmm ..."    
        t=ThreadRun(sw.RunSwmmDll,sp,rptfile,binfile)  
        t.start()
        t.join()
        ret=t.ret
        if(ret>0):
            print ret, "Error!" #sw.ENgeterror(e,25) 
            return None
    
        print "Success"
        f = swmmout.open(binfile)
        
        types=['subcatchments', 'nodes', 'links', 'system']
        ids={}
        times=[]
        for item in types:
            v=f.get_values(item)
            ids[item]=map( lambda x: x[0], v[0][1:])
            times=map( lambda x: x[0],v)
        self.ids=ids
        self.times=times
        
    def readCalibFile(self):
        try: 
            with open(self.parameters.calfile,'r') as f:
                lines = f.read().splitlines()   
            lines=[s.strip() for s in lines if not s.strip()[0]==';']
            if not lines[0]==self.parameters.calid[1]:
                print "Problem: The id specified on calibration file, ", lines[0], " is different from ", self.parameters.calid[1]
                raise
            if not lines[1]==self.parameters.caltype[2]:
                print "Problem: The variable type specified in file , ", lines[1], " is different from ", self.parameters.caltype[2]
                raise 
            id_=lines[0]
            type_=lines[1]
            lines=map(lambda x: float(x), lines[2:])
            if not len(lines)== len(self.times):
                print "problem: the length of data in the file, ", len(lines), " is different from expected: ", len(self.times)
                raise
            
            
        except: 
            print "Problem reading calibration file: ", self.parameters.calfile
            return None
        # now write a calibration file for swmm
        datf=self.parameters.calfile+".DAT"
        try:
            l=[self.times,lines]
            data=map(list,zip(*l))

            with open(datf,'w') as f:
                f.write("; Written by "+os.path.basename(__file__)+"\n")
                f.write(";"+self.parameters.caltype[2]+"\n")
                f.write(self.parameters.calid[1]+"\n")

                for time,value in data:
                    f.write(time.strftime('\t%m/%d/%Y\t%H:%M') + "\t" + str(value) + "\n")
            # now write a small ini file to be used as a template for all resulting inp files (Best_of_gen_xyz.inp)
            inif=self.parameters.projectdirectory+os.sep+"TEMPLATE.INI"
            with open(inif,'w') as f:
                f.write("[Calibration]\nFile%i=%s\n" %(swmm_ea_controller.SWMMCALIBRATIONFILE[self.parameters.caltype[0]], datf))
            self.parameters.calINITEMPLATE=inif
        except:
            print "Problem writing calibration file for swmm: ", datf
            return None            
                
        return lines    
            
            

    def initialize_optimization(self):
        print "Initializing the optimization ..."
        print "*****************************************************************"
        print "List of parameters:"
        print self.parameters
        print "*****************************************************************"
        parameters=self.parameters
        parameters.bestlist=[]
        for i in range(parameters.pop_size+1):
            parameters.bestlist.append([])          
        parameters.datadirectory = "." 
        parameters.resultsdirectory= "output"        
        parameters.projectdirectory=self.dirname
        parameters.templatefile=self.slotted_swmmfilename
        if parameters.swmmouttype[0]==swmm_ea_controller.SWMMREULTSTYPE_CALIB:
            parameters.calibdata=self.readCalibFile()
            parameters.swmmResultCodes=[parameters.caltype[1][0], parameters.calid[0],parameters.caltype[1][1]]
            if not parameters.calibdata: 
                return False
        self.swmm5ec=swmm5ec.SwmmEA()
        self.swmm5ec.setParams(parameters)
        self.swmm5ec.initialize()
        return True


    def write_slotted_swmm_file(self, full_fname, text):
        f=open(full_fname,'w')
        f.write(text)
        f.close()      
        self.slotted_swmmfilename=os.path.basename(str(QtCore.QDir.toNativeSeparators(full_fname)))

    def setswmmfile(self,swmmfilename):
    # Sets the swmm file swmmfilename that is in the project directory as the current one and Loads it. 
        self.setSwmmfile(swmmfilename)
        try:
            self.read_in_swmm_file()
            return True
        except: 
            self.setSwmmfile(None)
            return False

    def getSlottedData(self):
        sf=self.dirname+os.sep+(self.slotted_swmmfilename or "")
        slot=True
        if not (self.slotted_swmmfilename and os.path.exists(sf)):
            # then we read the original swmm file
            slot=False
            sf=self.dirname+os.sep+self.swmmfilename
        f=open(sf)
        t=f.read()
        f.close()
        if (not slot) and self.parameters.swmmouttype[0]==swmm_ea_controller.SWMMREULTSTYPE_STAGE: # we've read original swmm file and this analysis is multiple
            return reduce(lambda x,y: x+y, (swmm_ea_controller.SWMMSTAGESEPERATOR % (i) + t for i in range(self.parameters.stages)))
        return t

    def load(self, dirname=None, swmmfilename=None):
        if not dirname:
            dirname=self.dirname
        # first check directory
        if not os.path.isdir(dirname):
            return False
        pfilename=dirname+os.sep+self.paramfilename
        if not os.path.exists(pfilename):
            return False
        print "Using parameter file :" , pfilename
        f = open(pfilename)   
        flag=self.string_to_param(f.read())
        if not flag: 
            print "Problem opening and loading : ", pfilename
            raise()
        
        f.close() 
        if not swmmfilename:
            import glob
            try:
                self.setSwmmfile(os.path.basename(glob.glob(dirname+os.sep+'*.inp')[0]))
            except:
                try:
                    self.setSwmmfile(os.path.basename(glob.glob(dirname+os.sep+'*.INP')[0]))
                except:
                    print "problem: no swmm files in the directory."
                    return self.LOAD_NOSWMMFILE
        else:
            self.setSwmmfile(swmmfilename)
        self.read_in_swmm_file()
        return self.LOAD_COMPLETE

    def read_in_swmm_file(self):
        f = open(self.dirname+os.sep+self.swmmfilename)
        self.swmm_data=f.read()
        f.close()
        sf=self.dirname+os.sep+(self.swmmfilename or "")+"_"
        if os.path.exists(sf):
            self.slotted_swmmfilename=os.path.basename(str(sf))
        else:
            self.slotted_swmmfilename=None


    def string_to_param(self, string ):
        """As a policy we ignore entries !!python/object: and read as dictionary"""
        st=string.replace("!!python/object:","#!!python/object:")
        f=StringIO.StringIO(st)
        self.parameters=parameters_class_()
        # now add the parameter to ensure compatibility with old version
        self.parameters.swmmouttype=[swmm_ea_controller.SWMMREULTSTYPE_FLOOD, swmm_ea_controller.SWMMCHOICES[swmm_ea_controller.SWMMREULTSTYPE_FLOOD]]# default
        self.parameters.stage_size=1
        self.parameters.discount_rate=1.0
        self.parameters.stages=1
        try:
            import yaml
            dataMap = yaml.load(f)  
            f.close()
            ##if the yaml file has the header !!python/object:swmmeaproject.parameters_class_
            ## yaml.load will load the object directly. 
            ## if not loads a dictionary
            #if(dataMap.__class__.__name__==self.parameters.__class__.__name__):
                #self.parameters=dataMap
            #else:
                #for key in dataMap :
                    #setattr(self.parameters, key, dataMap[key])
                    
            for key in dataMap :
                setattr(self.parameters, key, dataMap[key])             
        except: 
            print "Problem reading parameters ",  sys.exc_info()
            return None
        return True        



    def cleanup_and_make_to_dict(self,params):
        """ 1. removes the temporary parameters in params object.
            2. convert it to a dictionary. 
            """
        from  copy  import deepcopy
        p=deepcopy(params)
        for item in [ "calINITEMPLATE", "calibdata", "bestlist", "linestring",  "projectdirectory",  "resultsdirectory", "templatefile", "datadirectory" ]: 
            try: 
                delattr(p,item)
            except: 
                pass
        return p.__dict__

    def save(self):

        print "Saving to directory : " + self.dirname
        if(self.swmmfilename):
            f=open(self.dirname+os.sep+self.swmmfilename,'w')
            f.write(self.swmm_data)
            f.close()
            print self.swmmfilename + " saved."
        f=open(self.dirname+os.sep+self.paramfilename,'w')
        import yaml
        f.write("#Written programmetically!\n")
        d=yaml.dump(self.cleanup_and_make_to_dict(self.parameters),f)
        f.close()

        print self.paramfilename + " saved."
        f.close
        self.makeemptydirectory(self.dirname+os.sep+self.outdir)
        self.makeemptydirectory(self.dirname+os.sep+self.tmpdir)



    def copy(self,dirname=None):
        if os.path.exists(dirname+os.sep+"param.yaml"):
            print "Please delete %s directory and try again!" % (dirname)
            return None
        try:
            self.save()
            copy_tree(str(QtCore.QDir.toNativeSeparators(self.dirname)),str(QtCore.QDir.toNativeSeparators(dirname)))
            # above copy_tree is from distutils
            return Project(dirname,self.swmmfilename)
        except:
            print "Error in creating new project location: "+ dirname
            return None

    def makeemptydirectory(self, out):
        if not os.path.isdir(out):
            if os.path.exists(out):
                os.remove(out)
            os.mkdir(out)
            import threading

import threading           
class ThreadRun(threading.Thread):
    def __init__(self, target, *args):
        self._target = target
        self._args = args
        threading.Thread.__init__(self)
 
    def run(self):
        self.ret=self._target(*self._args)

if __name__ == "__main__":
    pr=Project()
    print pr.dirname
    print pr.paramfilename
    print pr.swmmfilename
    pr.save()
    npr=pr.copy("C:\\Users\\apa.EDUCATION\\Desktop\\tmpswmmea")
    if not npr:
        print "copying failed!"
    else:
        print npr.dirname
        print npr.paramfilename
        print npr.swmmfilename 
        pr2=Project("C:\\Users\\apa.EDUCATION\\Desktop\\tmpswmmea")
        print "hi"