import rngStream
from scipy.stats import norm # normal distribution for CIs and generation
from argparse import ArgumentParser
from math import sqrt, log, pi, cos, sin
from copy import copy

# This Program Simulates a GIG1 queue under various input processes
#
# CS 590M: Simulation
#
# Usage: python GIG1.py -n k [-t -d m -p a]
# where -n (--num) k: Number of replications given by k
# -t (--trial): Trial run to get required number of replications
# -d (--debug) m: Verbose output needed, levels from 0-2
# -p (--ptype) a: input process model
#    a=0 - Poisson process with mean mu
#    a=1 - Weibull distribution, mean mu, SD = 0.5 * mu
#    a=2 - Weibull distribution, mean mu, SD = 2.0 * mu
#    a=3 - Proc. w. neg. correlated mean mu exponential increments via ARTA
#    a=4 - Proc. w. pos. correlated mean mu exponential increments via ARTA
#
# Uses the rngStream module (the module should be in the Python path or
# the same folder)
__author__ = "Peter Haas"


class Estimator:
    """ Computes point estimates and confidence intervals """
    def __init__(self, conf_pct, epsilon, relative):
        self.epsilon = epsilon # precison of CI
        self.relative = relative # relative or absolute precision
        self.k = 0  # number of values processed so far
        self.sum = 0.0  # running sum of values
        self.v = 0.0  # running value of (k-1)*variance
        self.z = norm.ppf(0.5 + (conf_pct / 200))  # CI normal quantile
        self.conf_str = str(conf_pct) + "%"  # string of form "xx%" for xx% CI
        
    def reset(self):
        self.k = 0
        self.sum = 0
        self.v = 0

    def process_next_val(self, value):
        self.k += 1
        if self.k > 1:
            diff = self.sum - (self.k - 1) * value
            self.v += diff/self.k * diff/(self.k-1)
        self.sum += value

    def get_variance(self):
        return self.v/(self.k-1) if self.k > 1 else 0

    def get_mean(self):
        return self.sum/self.k if self.k > 1 else 0

    def get_conf_interval(self):
        hw = self.z * sqrt(self.get_variance()/self.k)
        point_est = self.get_mean()
        c_low = point_est - hw
        c_high = point_est + hw
        rstr = "{0} Confidence Interval: [ {1:.4f}, {2:.4f} ]  (hw = {3:.4f})"
        return rstr.format(self.conf_str, c_low, c_high, hw)

    def get_num_trials(self):
        var = self.get_variance()
        if var == 0:
            return "UNAVAILABLE (Need at least 2 pilot reps)"
        if self.relative:
            width = self.get_mean() * self.epsilon
            rstr = "Est. # of reps for +/- {0:.2f}% accuracy: {1}"
            acc = 100 * self.epsilon
        else:
            width = self.epsilon
            rstr = "Est. # of reps for +/- {0:.4f} accuracy: {1}"
            acc = self.epsilon
        nreps = int(1+(var * self.z * self.z)/(width * width))

        return rstr.format(acc, nreps)


def verbose_print(level, *args):
    """ verbose printing for different debug levels """
    if level <= int(sysargs.debug):
        for arg in args:
            print(arg, end=' ')
    else:
        pass


class Sim_Model:
    """ Instantiates GI/G/1 queue model """
    def __init__(self, p_type, trial):
        self.p_type = p_type  # performance measure to compute
        self.num_ptypes = 5  #  total number of performance measures
        self.perf_name = ""  # name of performance measure
        self.num_gens = 9  # number of rng streams needed
        self.m_serv = 0.99  # mean service time
        self.m_inter = 1.0   # mean interarrival time
        self.wlam0 = 0.8856899  #  Weibull params for mu = 1, sd = 0.5
        self.walpha0 = 2.1013491
        self.wlam1 = 1.7383757  #  Weibull params for mu = 1, sd = 2.0
        self.walpha1 = 0.5426926
        self.cp = 1.0/sqrt(2.0)  # correlation parameter
        self.sim_length = 500  # length of simulation
        self.unigens = {} # initialize dict of generators
        self.num_normals = 100  # number of normals to store
        self.normals = []  # array of normal random variables
        self.last_normal = 0  # previously generated normal
        self.unigens["exp"] = rngStream.RngStream()
        self.unigens["serv0"] = rngStream.RngStream()
        self.unigens["serv1"] = rngStream.RngStream()
        self.unigens["weib0"] = rngStream.RngStream()
        self.unigens["weib1"] = rngStream.RngStream()
        self.unigens["norm0_0"] = rngStream.RngStream()
        self.unigens["norm0_1"] = rngStream.RngStream()
        self.unigens["norm1_0"] = rngStream.RngStream()
        self.unigens["norm1_1"] = rngStream.RngStream()
        # use different streams for trial runs
        if trial:
            for key in self.unigens.keys():
                for j in range(self.num_ptypes):
                    self.unigens[key].ResetNextSubstream()
        # choose different rng stream for different performance measures
        for key in self.unigens.keys():
            for j in range(int(p_type)):
                self.unigens[key].ResetNextSubstream()
        # store name of performance measure
        if p_type == 0:
            self.perf_name = "Poisson"
        elif p_type == 1:
            self.perf_name = "Weibull, sd = 0.5 * mu"
        elif p_type == 2:
            self.perf_name = "Weibull, sd = 2.0 * mu"
        elif p_type == 3:
            self.perf_name = "Neg. correlated Poisson"
        else:
            self.perf_name = "Pos. correlated Poisson"
        # Initialize self.normals
        if self.p_type >= 3:
            self.genNormals(True)
        
    def genService(self):
        """ Generate service time via triangular distn """
        u1 = self.unigens["serv0"].RandU01()
        u2 = self.unigens["serv1"].RandU01()
        x = (u1+u2) * self.m_serv
        return x
    
    def genNormals(self, init):
        """ fills up self.normal with num_normals N(0,1) variates """
        """ uses rng generators index by gkey0 and gkey1 """
        """ If init is true, then initializes self.last_normal """
        gkey0 = "norm0_0" if self.p_type <= 3 else "norm1_0"
        gkey1 = "norm0_1" if self.p_type <= 3 else "norm1_1"
        i = 1
        while i < self.num_normals:
            u1 = self.unigens[gkey0].RandU01()
            u2 = self.unigens[gkey1].RandU01()
            r = sqrt(-2.0 * log(u1))
            theta =  2 * pi * u2;
            self.normals.append(r * cos(theta))
            self.normals.append(r * sin(theta))
            i += 2
        if init:
            self.last_normal = self.normals.pop()

    def genInter(self):
        """ generate interarrival time depending on p_type """
        # refill vector of normals if needed
        if self.p_type >=3 and len(self.normals) == 0:
            self.genNormals(False)  # no initialization
            
        # generate interarrival time
        if self.p_type == 0:  #exponential
            u1 = self.unigens["exp"].RandU01()
            val = -self.m_inter * log(u1)
        elif self.p_type == 1:  # Weibull with SD = 0.5 mu
            u1 = self.unigens["weib0"].RandU01()
            val = pow(-log(u1), 1.0 / self.walpha0) / self.wlam0
        elif self.p_type == 2: # Weibull with SD = 2.0 mu
            u1 = self.unigens["weib1"].RandU01()
            val = pow(-log(u1), 1.0 / self.walpha1) / self.wlam1
        elif self.p_type == 3: # neg-correlated exponential
            z = self.normals.pop()
            val = self.cp * (z - self.last_normal)  # val is neg.-corr. normal
            val = norm.cdf(val)  # val is now neg.-corr. uniform
            val = -self.m_inter * log (val)  # val is now neg.-corr. exp.
            self.last_normal = z
        else: # pos-correlated exponential
            z = self.normals.pop()
            val = self.cp * (z + self.last_normal)  # val is neg.-corr. normal
            val = norm.cdf(val)  # val is now neg.-corr. uniform
            val = -self.m_inter * log (val)  # val is now neg.-corr. exp.
            self.last_normal = z
        return val
    
    def print_sc(self, state, clocks, time, istar):
        """ Format state, clock readings, and trigger event for printing """
        cf = "({:7.4f}, {:7.4f})"  # format string for clock readings
        sf = "({:4d})"  #format string for state
        cstr = cf.format(clocks[0], clocks[1])
        sstr = sf.format(state)
        istr = "i* = {:1d}".format(istar) if istar >= 0 else ""
        tstr = "T = {0:8.4f}".format(time)
        outstr = "{}: {}, {} {}\n".format(tstr, sstr, cstr, istr)
        return outstr
        
    def do_rep(self, est):
        """ Computes specified performance meas. for a single replication """
        state = 1 # initialize state
        clocks = [-1, -1] # initialize clock-reading vector
        time = 0.0
        done = False  # true if simulation replication has completed
        perf_meas = 0  # initialize performance measure
        # schedule initial clock readings
        clocks[0] = self.genInter()
        clocks[1] = self.genService()
        verbose_print(2, "\n")
        while not done:  # generate the replication
        
            # compute holding time and trigger event
            htime = min(t for t in clocks if t > 0)
            istar = clocks.index(htime)
            
            # set (s',c') = (s,c) as default
            clocks_next = copy(clocks)  # ensures values are copied
            state_next = copy(state)
            
            # compute updated clock reading vector (will update it more below)
            clocks_next = [x - htime if x > 0 else x for x in clocks]
            clocks_next[istar] = -1  # clock for trigger event has run down

            # compute next state s' and next clocks c' (i.e., schedule events)
            # branch according to which event triggers next transition
            if istar == 0:  # arrival
                state_next += 1
                clocks_next[0] = self.genInter()
                if state == 0:  # arrival to empty queue
                    clocks_next[1] = self.genService()
            else: # service completion
                state_next -= 1
                if state_next > 0:  # schedule next service completion
                    clocks_next[1] = self.genService()
            # end of (s', c') computation
            
            #  debug print of (s, c, e*)
            verbose_print(2, self.print_sc(state, clocks, time, istar))
           
            # compute/update performance measure and check for termination
            if time + htime > self.sim_length:  # sim. rep is finished
                done = True
                htime = self.sim_length - time
            perf_meas += state * htime
            if done:  # finish performance measure computation
                perf_meas /= self.sim_length
                    
            # update state and time
            state = copy(state_next) 
            clocks = copy(clocks_next)
            time += htime
        # End of while loop (end of simulation repetition)
        
        # final detailed debug print
        verbose_print(2, self.print_sc(state, clocks, time, -1))
            
        # update statistics and return
        est.process_next_val(perf_meas)
        return perf_meas

    
if __name__ == "__main__":
    # parse command line arguments
    parser = ArgumentParser(
        description="mancell -num n [--trial --debug m --ptype a]")
    parser.add_argument('-n', '--num', help="Number of replications",
                        required=True)
    parser.add_argument('-t', '--trial',
                        help="Trial run to get required # of replications",
                        action='store_true')
    parser.add_argument('-d', '--debug', help="Set verbose output level",
                        default=0)
    parser.add_argument('-p', '--p_type',
                        help="Performance measure: 0, 1, 2, or 3",
                        default=0)
    sysargs = parser.parse_args()
    
    # ensure that results are reproducible
    seed0 = [12345, 12345, 12345, 12345, 12345, 12345]
    rngStream.SetPackageSeed(seed0)

    # instantiate estimator
    est = Estimator(95, 0.01, True)  # 95% CI with 1% relative precision

    # Instantiate simulation model
    sim = Sim_Model(int(sysargs.p_type), sysargs.trial)
    
    # run simulation repetitions and collect stats
    #print("\n")
    for rep in range(int(sysargs.num)):
        val = sim.do_rep(est)
        verbose_print(1, "\nRepetition {0:3d}: {1:.2f}\n".format(rep+1, val))
        #if rep % 1000 == 0: print(rep)

    # print results
    verbose_print(0, "\n" + sim.perf_name)
    verbose_print(0, "\nAvg. queue length: {:.2f}".format(est.get_mean()))
    verbose_print(0, "\n{}\n".format(est.get_conf_interval()))
    if sysargs.trial:
        verbose_print(0, "{}".format(est.get_num_trials()))