import numpy as np
np.seterr(all='raise')
import pickle
import os
import time
from minimizer import Run_Result,tracing

data_path = "/home/dish/varmint/data/"

ksq = 9

def readable_time(seconds):
    seconds = int(np.round(np.abs(seconds)))
    minutes = seconds//60
    seconds = seconds%60
    hours = minutes//60
    minutes = minutes%60
    days = hours//24
    hours = hours%24
    if days>0:
        return "{0:,}d:{1}h".format(days,hours)
    if hours>0:
        return "{0}h:{1}m".format(hours,minutes)
    if minutes>0:
        return "{0}m:{1}s".format(minutes,seconds)
    return "{0}s".format(seconds)

class Container(object):
    """
    a container for all of the data generated by and needed for running varmint
    Provides methods to update itself after runs and check convergence
    """
    @tracing
    def __init__(self,N,u,Tstep,numT,runs_req,max_bad_runs):
        self.n = N
        self.u = u

        self.Tstep = Tstep
        self.numT = numT       


        self.runs_req = runs_req
        self.max_bad_runs = max_bad_runs


        self.time     = 0
        self.time2    = 0
        self.num_runs = 0

        self.EV  = 0
        self.EV2 = 0
        self.EV4 = 0


        self.F   = np.zeros(self.numT)
        self.F2  = np.zeros(self.numT)
        self.F4  = np.zeros(self.numT)
        self.S   = np.zeros(self.numT)
        self.S2  = np.zeros(self.numT)
        self.S4  = np.zeros(self.numT)
        self.FS  = np.zeros(self.numT)
        self.FS2 = np.zeros(self.numT)
        self.CV  = np.zeros(self.numT)
        self.CV2 = np.zeros(self.numT)
        self.CV4 = np.zeros(self.numT)

        self.C2  = np.zeros(self.numT)
        self.C4  = np.zeros(self.numT)
        self.C8  = np.zeros(self.numT)
        self.X   = np.zeros(self.numT)
        self.X2  = np.zeros(self.numT)

        self.G2 = np.zeros((self.numT,self.runs_req))

        
        self.num_bad_runs = 0
        self.bad_ev = np.array([])
        self.err_collection = {}

    @tracing
    def set_stop_cond(self,runs_req,max_bad_runs):
        if self.runs_req > runs_req:
            self.G2 = self.G2[:,:max(self.num_runs,runs_req)]
        elif self.runs_req < runs_req:
            self.G2 = np.pad(self.G2,((0,0),(0,runs_req-self.runs_req)),'constant')
        self.runs_req = runs_req
        self.max_bad_runs = max_bad_runs

    @tracing
    def more_runs(self,quit_time):

        if self.num_runs >= self.runs_req:
            return False

        if self.num_bad_runs>=self.max_bad_runs:
            return False     

        return True

    @tracing
    def record(self,res):
        if res.stat == 's':
            print "recording run"
            self.num_runs = self.num_runs+1

            self.time = self.time+(res.elapsed - self.time)/self.num_runs
            self.time2 = self.time2+(res.elapsed**2 - self.time2)/self.num_runs

            self.EV = self.EV+(res.ev - self.EV)/self.num_runs
            self.EV2 = self.EV2+(res.ev**2 - self.EV2)/self.num_runs
            self.EV4 = self.EV4+(res.ev**4 - self.EV4)/self.num_runs

            self.F = self.F+(res.fe - self.F)/self.num_runs
            self.F2 = self.F2+(res.fe**2 - self.F2)/self.num_runs
            self.F4 = self.F4+(res.fe**4 - self.F4)/self.num_runs

            self.S = self.S+(res.s - self.S)/self.num_runs
            self.S2 = self.S2+(res.s**2 - self.S2)/self.num_runs
            self.S4 = self.S4+(res.s**4 - self.S4)/self.num_runs

            self.FS = self.FS+(res.fe*res.s - self.FS)/self.num_runs
            self.FS2 = self.FS2+((res.fe*res.s)**2 - self.FS2)/self.num_runs

            self.CV = self.CV+(res.cv - self.CV)/self.num_runs
            self.CV2 = self.CV2+(res.cv**2 - self.CV2)/self.num_runs  
            self.CV4 = self.CV4+(res.cv**4 - self.CV4)/self.num_runs  

            self.C2 = self.C2+ (res.c2 - self.C2)/self.num_runs
            self.C4 = self.C4+ (res.c2**2 - self.C4)/self.num_runs
            self.C8 = self.C8+ (res.c2**4 - self.C8)/self.num_runs

            if self.num_runs<=self.runs_req:
                self.G2[:,self.num_runs-1] = res.g2
            else:
                res.g2.shape = (res.g2.shape[0],1)
                self.G2 = np.concatenate((self.G2,res.g2),axis=1)


            self.X = self.X+ (res.x - self.X)/self.num_runs
            self.X2 = self.X2+ (res.x**2 - self.X2)/self.num_runs

        else:
            self.num_bad_runs = self.num_bad_runs+1
            self.bad_ev = np.append(self.bad_ev,res.ev)
            code = res.stat.split()
            self.record_err((int(code[1]),int(code[2])),res.mess)


    @tracing
    def record_err(self,code,message):

        self.num_bad_runs = self.num_bad_runs+1
        if message in self.err_collection.keys():
            if code in self.err_collection[message].keys():
                self.err_collection[message][code] = self.err_collection[message][code]+1
            else:
                self.err_collection[message][code] = 1
        else:
            self.err_collection[message] = {code:1}

    @tracing
    def save_self(self,flag):
        if not flag == '0' and not flag == '1':
            raise ValueError('Invalid Flag')
        if flag == '0':
            self.G2 = self.G2[:,:self.num_runs]
        save_file = open(data_path+"{0}_{1}_temp.d".format(self.n,self.u),'w')
        save_file.write(flag+"\n")
        pickle.dump(self,save_file)
        save_file.close()
        os.rename(data_path+"{0}_{1}_temp.d".format(self.n,self.u),
                    data_path+"{0}_{1}.d".format(self.n,self.u))

    @tracing
    def dump_self(self):
        save_file = open(data_path+"{0}_{1}_dump.d".format(self.n,self.u),'w')
        pickle.dump(self,save_file)
        save_file.close()

    @tracing
    def err_report(self):
        report = "Errors encountered:\n\n"
        for err in self.err_collection.keys():
            report = report+err+"\n"
            report = report+"in: "+str(self.err_collection[err])+"\n\n"

        return report

    @tracing
    def status(self):
        message = "N = {0} u = {1}\n".format(self.n,self.u)
        message = message+"T = {0} to T = {1} in steps of {2}\n".format(self.Tstep,self.Tstep*self.numT,self.Tstep)
        message = (message
            +"Requested Runs = {0:,},".format(self.runs_req)
            +" Bad runs = {0:,} of {1:,}\n\n".format(self.num_bad_runs,self.max_bad_runs))
        if self.num_runs<=0:
            message = message+"No runs yet completed\n"
        elif self.num_runs == 1:
            message = message+"One run complete\nCompute time = "
            message = message+readable_time(self.time)+"\n"
            message = message+"Estimated compute time to requested runs = "
            message = message+readable_time((self.runs_req-self.num_runs)*self.time)+"\n"
        elif self.num_runs < self.runs_req:
            timevar = 3*np.sqrt(1.0+1.0/(self.num_runs-1))*np.sqrt(self.time2-self.time**2)
            message = message+"{0:,} runs completed\n".format(self.num_runs)
            message = message+"Compute time per run = "
            message = message+readable_time(self.time)+" +/- "+readable_time(timevar)+"\n"
            message = message+"Estimated compute time to requested runs = "
            message = message+readable_time((self.runs_req-self.num_runs)*self.time)+" +/- "
            message = message+readable_time((self.runs_req-self.num_runs)*timevar)+"\n"
        else:
            timevar = 3*np.sqrt(1.0+1.0/(self.num_runs-1))*np.sqrt(self.time2-self.time**2)
            message = message+"{0:,} runs completed\nCompute time per run = ".format(self.num_runs)
            message = message+readable_time(self.time)+" +/- "+readable_time(timevar)+"\n"
            message = message+"All requested runs completed.\n"


        message = message+"\n"+self.err_report()
        return message

@tracing
def open_data(N,u):

    save_file = open(data_path+"{0}_{1}.d".format(N,u),'r+')
    flag = save_file.readline()
    if flag == "1\n":
        save_file.close()
        raise Exception('open')
    elif flag == "0\n":
        save_file.seek(0)
        save_file.write('1')
        save_file.seek(2)
        try:
            bundle = pickle.load(save_file)
        except Exception as e:   
            save_file.seek(0)
            save_file.write('0')
            save_file.close()
            raise e

        save_file.close()
    else:
        save_file.close()
        raise Exception('flag')

    if bundle.G2.shape[0]<bundle.runs_req:
        bundle.G2 = np.pad(bundle.G2,((0,0),(0,bundle.runs_req-bundle.G2.shape[0])),'constant')
    return bundle

