import rngStream
import numpy as np
import heapq  # For event list management
from scipy.stats import norm, t  # normal and Student-t distributions for CIs
from argparse import ArgumentParser
from math import sqrt, log, ceil
from copy import copy

# Simulates the telephone model of Assignment #7
#
# CS 590M: Simulation
#
# Usage: python phone.py [-n k -t -d m -p a]
# where -n (--num) k: Number of replications given by k [default 500)
# -t (--trial): Trial run to get required number of replications
# -d (--debug) m: Verbose output needed, levels from 0-2
# -r (--rate) specifies the phone rate (default 5 cents per minute)
# -p (--ptype) a: performance measure desired
#    a=0 - IQR for amount of time in [0,100] that all links are idle,
#    a=1 - long-run frac. of time that link 2 is busy
#    a=2 - stddev of amount of time in [0,100] that all links are idle
#    a=3 - expected total revenue in [0,100]
#
# Uses the rngStream module (the module should be in the Python path or
# the same folder)
__author__ = "Peter Haas"

class EventList:
    """Event list using a heap data structure"""
    def __init__(self):
        self.event_list = []
        self.event_record = {}

    def add_event(self, event_id, time):
        """Adds event_id with time in the event list"""
        if event_id in self.event_record:
            self.cancel_event(event_id)
        new_event = [time, event_id]
        self.event_record[event_id] = new_event
        heapq.heappush(self.event_list, new_event)

    def cancel_event(self, event_id):
        """Cancels event having id = event_id"""
        if event_id in self.event_record:
            to_remove = self.event_record.pop(event_id)
            to_remove[-1] = "<canceled>"
        else:
            raise KeyError("Event %s not in list." %str(event_id))

    def next_event(self):
        """Return tuple containing (event_id, time)"""
        if self.event_list:
            next_event = heapq.heappop(self.event_list)
            if next_event[-1] != "<canceled>":
                del self.event_record[next_event[-1]]
                return next_event[-1], next_event[0]
            return self.next_event()
        else:
            raise KeyError("Popping from an empty event list.")

    def reset(self):
        self.event_list = []
        self.event_record = {}

    def __str__(self):
        el = sorted(self.event_list)
        return "\n".join([str(x) for x in el])


class Estimator:
    """ Computes point estimates and confidence intervals """
    def __init__(self, conf_pct, epsilon, relative):
        self.epsilon = epsilon  # precison of CI
        self.conf_pct = conf_pct  # confidence of CI as a percent
        self.relative = relative  # relative or absolute precision
        self.point_est = 0  # point estimate
        self.cihw = 0  # CI half-width
        self.svar = 0
        self.data = []  # holds simulation observations
#        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.data = []
        self.pt_est = 0
        self.se_est = 0
        self.svar = 0

    def get_next_val(self, value):
        self.data.append(value)
        
    def IQR_est_jack(self, num_sections):
        """ Estimates interquartile range using jackknife + sectioning """
        #  Compute est quantile over all data
        num_obs = len(self.data)
        lo = ceil(0.25 * num_obs) - 1
        hi = ceil(0.75 * num_obs) - 1
        obs = copy(self.data)
        obs.sort()
        alpha_n = obs[hi] - obs[lo]        
        # Compute pseudovalues
        section_length = int(num_obs / num_sections)
        lo = ceil(0.25 * (num_obs - section_length)) - 1
        hi = ceil(0.75 * (num_obs - section_length)) - 1        
        pseudo = []
        for i in range(num_sections):
            obs = copy(self.data)
            # delete section:
            obs[(section_length * i):(section_length * (i + 1))] = []
            obs.sort()
            iqr = obs[hi] - obs[lo]
            pseudo.append(num_sections * alpha_n - (num_sections - 1) * iqr)
        self.point_est = np.mean(pseudo)
        std_error = np.std(pseudo, ddof=1) / sqrt(num_sections)
        t_quantile = t.ppf(0.5 + (self.conf_pct / 200), df = num_sections - 1)
        self.cihw = t_quantile * std_error
        
    def est_regen(self):
        """ computes regenerative point est. and CI from (Y, tau) pairs """
        Y = [x[0] for x in self.data]
        tau = [x[1] for x in self.data]
        num_obs = len(self.data)
        s = np.cov(Y, tau, ddof=1)
        r = sum(Y) / sum(tau)
        taubar = np.mean(tau)
        sigma2 = (s[0,0] - 2 * r * s[0,1] + r * r * s[1,1]) / (taubar * taubar)
        self.point_est = r
        z = norm.ppf(0.5 + (self.conf_pct / 200))  # CI normal quantile
        self.cihw = z * sqrt(sigma2 / num_obs)
        
    def STD_est_Taylor(self):
        """ estimates std dev using Taylor-series method """
        num_obs = len(self.data)
        avgx = np.mean(self.data)
        avgx2 = np.mean([x**2 for x in self.data])
        cn = -avgx / sqrt(avgx2 - avgx * avgx)
        dn = 0.5 / sqrt(avgx2 - avgx * avgx)
        sn2 = 0
        for i in range(num_obs):
            x = self.data[i]
            temp = cn * (x - avgx) + dn * (x * x - avgx2)
            sn2 += temp * temp
        sn2 /= (num_obs - 1)
        self.point_est = sqrt(avgx2 - avgx * avgx)
        z = norm.ppf(0.5 + (self.conf_pct / 200))  # CI normal quantile
        self.cihw = z * sqrt(sn2 / num_obs)
        
    def mean_est(self):
        num_obs = len(self.data)
        self.point_est = np.mean(self.data)
        self.svar = np.var(self.data, ddof=1)
        std_error = sqrt(self.svar / num_obs)
        z = norm.ppf(0.5 + (self.conf_pct / 200))  # CI normal quantile
        self.cihw = z * std_error

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

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 phone model """
    def __init__(self, p_type, trial, rate):
        self.p_type = p_type  # performance measure to compute
        self.num_ptypes = 4  # total number of performance measures
        self.num_lines = 4
        self.num_links = 2
        self.phone_rate = rate
        self.max_length = 6.0  # max length of call (mins)
        self.mean_wait = 6.0  # mean length of time until call placed (mins)
        if rate > 0.05:
             self. mean_wait *= pow(rate/0.05,1.3)
        self.sim_length = 100  # length of simulation (for p_type = 1, 3, 4)
        # set up random number streams
        line_uni = []  # holds streams for placement of calls
        link_uni = []
        for i in range(self.num_lines):
            line_uni.append(rngStream.RngStream())  # streams for placing calls
        for i in range(self.num_links):
            link_uni.append(rngStream.RngStream())  # streams for call lengths
        callID = rngStream.RngStream()  # stream for choosing called line
        self.unigens = {"line":line_uni, "link":link_uni, "callID":callID}
        # use different streams for trial runs
        if trial:
            for i in range(self.num_lines):
                for j in range(self.num_ptypes):
                    self.unigens["line"][i].ResetNextSubstream()
            for i in range(self.num_links):
                for j in range(self.num_ptypes):
                    self.unigens["link"][i].ResetNextSubstream()
            for j in range(self.num_ptypes):
                self.unigens["callID"].ResetNextSubstream()
        # choose different rng stream for different performance measures
        for i in range(self.num_lines):
            for j in range(int(self.p_type)):
                self.unigens["line"][i].ResetNextSubstream()
        for i in range(self.num_links):
            for j in range(int(self.p_type)):
                self.unigens["link"][i].ResetNextSubstream()
        for j in range(int(self.p_type)):
            self.unigens["callID"].ResetNextSubstream()

    def print_se(self, state, elist, time):
        """ Format state, trigger event, and event list for printing """
        sstr = '(' + ','.join('{:2d}'.format(x) for x in state) +')'
        #istr = "i* = {:1d}".format(istar) if istar >= 0 else ""
        tstr = "T = {0:8.4f}".format(time)
        outstr = "{}: {}\n".format(tstr, sstr)
        outstr += str(elist) +"\n"
        return outstr
        
    def do_rep(self, est, elist):
        """ Computes specified performance meas. for a single replication """
        elist.reset()  # reset event list
        state = [0] * self.num_lines  # initialize state
        regen_state = [0] * self.num_lines  # regen = entrance to this state
        time = 0.0
        perf_meas = 0  # for all methods other than regenerative
        Y = 0  # for regenerative method
        tau = 0  # for regenerative method
        done = False  # true if simulation replication has completed
        # schedule initial events
        for l in range(self.num_lines):
            u = self.unigens["line"][l].RandU01()
            elist.add_event(l + 1, -self.mean_wait * log(u))
        verbose_print(2, "\n")
        while not done:  # generate the replication
        
            #  debug print of state and event list
            verbose_print(2, self.print_se(state, elist, time))

            # compute holding time and trigger event
            istar, e_time = elist.next_event()
            h_time = e_time - time

            # set next state = current state as default
            state_next = copy(state)
            
            # compute useful quantities from state
            busy_links = [l + 1 for l in range(self.num_links)\
              if l + 1 in state]
            num_busy_links = len(busy_links)
            all_busy = (num_busy_links == self.num_links)
            all_idle = (num_busy_links == 0)
            avail_link = min([l + 1 for l in range(self.num_links)\
              if l + 1 not in busy_links], default=None)                        

            # compute next state and schedule events as needed
            # branch according to which event triggers next transition
            if istar <= self.num_lines:  # trigger event is a placed call
                # determined called line (!= istar)
                u = self.unigens["callID"].RandU01()
                called_line = ceil(u * (self.num_lines - 1))
                if called_line >= istar:
                    called_line += 1
                verbose_print(2,"called line: {:2d}\n".format(called_line))
                if (not all_busy) & (state[called_line - 1] == 0):  #  success
                    state_next[istar - 1] = avail_link
                    state_next[called_line - 1] = avail_link
                    # schedule end of call
                    u = self.unigens["link"][avail_link - 1].RandU01()
                    elist.add_event(self.num_lines + avail_link,\
                      e_time + u * self.max_length)
                    # cancel placement of call by called line
                    elist.cancel_event(called_line)
                else:  # call fails
                    # schedule next call placement by line istar
                    u = self.unigens["line"][istar - 1].RandU01()
                    elist.add_event(istar, e_time - self.mean_wait * log(u))
            else:  # trigger event is a completion of call
                # set appropriate state components equal to 0
                finished_link = istar - self.num_lines
                for l in range(self.num_lines):
                    if state[l] == finished_link:
                        state_next[l] = 0
                        # schedule next call placement
                        u = self.unigens["line"][l].RandU01()
                        elist.add_event(l + 1,\
                          e_time - self.mean_wait * log(u))
            # end of state update and event scheduling
          
            # compute/update performance measure and check for termination
            if self.p_type in {0, 2}:  # amount of time all links are idle
                if time + h_time > self.sim_length:  # sim. rep is finished
                    h_time = self.sim_length - time
                    done = True
                if all_idle:
                    perf_meas += h_time
            elif self.p_type == 1:  # amount of time in cycle w. link 2 busy
                if state_next == regen_state:
                    done = True
                tau += h_time
                if 2 in busy_links:
                    Y += h_time
            else:  # p_type == 3, expected total revenue
                if time + h_time > self.sim_length:  # sim. rep is finished
                    h_time = self.sim_length - time
                    done = True
                perf_meas += num_busy_links * self.phone_rate * h_time

            # update state and time
            state = copy(state_next) 
            time += h_time
        # End of while loop (end of simulation repetition)
        
        # final detailed debug print
        verbose_print(2, self.print_se(state, elist, time))
            
        # update statistics and return
        if self.p_type == 1:  # regenerative estimation
            perf_meas = (Y, tau)
        est.get_next_val(perf_meas)
        return perf_meas

    
if __name__ == "__main__":
    # parse command line arguments
    parser = ArgumentParser(
        description="phone -num n [--trial --debug m --rate x --ptype a]")
    parser.add_argument('-n', '--num', help="Number of replications",
                        default = 500)
    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('-r', '--rate', help="Per-minute phone rate in $",
                        default=0.05)
    parser.add_argument('-p', '--p_type',
                        help="Performance measure: 0, 1, 2, or 3",
                        default=0)
    sysargs = parser.parse_args()
    
    num_sections = 5  # number of sections for jackknife IQR estimate
    
    # ensure that results are reproducible
    seed0 = [12345, 12345, 12345, 12345, 12345, 12345]
    rngStream.SetPackageSeed(seed0)

    # instantiate estimator
    est = Estimator(95, 0.01, True)  # 99% CI with 1% relative precision
    
    # instantiate event list
    elist = EventList()

    # Instantiate simulation model
    sim = Sim_Model(int(sysargs.p_type), sysargs.trial, float(sysargs.rate))
    
    # run simulation repetitions and collect stats
    for rep in range(int(sysargs.num)):
        val = sim.do_rep(est, elist)
        verbose_print(1, "\nRepetition {:3d}: {}\n".format(rep+1, str(val)))

    # compute estimates
    if sim.p_type == 0:
        pstr = "\nIQR[time in [0,100] that all links idle]: {:.2f}"
        est.IQR_est_jack(num_sections)
    elif sim.p_type == 1:
        pstr = "\nlong-run frac of time that link 2 is busy]: {:.4f}"
        est.est_regen()
    elif sim.p_type == 2:
        pstr = "\nStddev[time in [0,100] that all links idle]: {:.2f}"
        est.STD_est_Taylor()
    else:
        pstr = "\nE[revenue in [0,100]]: {:.2f}"
        est.mean_est()
        
    # print results
    verbose_print(0, pstr.format(est.point_est))
    verbose_print(0, "\n{}".format(est.get_conf_interval()))
    if sim.p_type == 3:
        verbose_print(0, "\nvar = {:.2f} ({:4d} reps.)".format(est.svar,\
                      int(sysargs.num)))
    print("\n")