import rngStream
from math import sqrt, log
from argparse import ArgumentParser

__author__ = "Peter Haas"

# This Program Simulates the hospital-unit GSMP model of Assigment #2
#
# CS 590M: Introduction to Simulation
#
# Usage: hospitalg -n k [-d m -t -p a]
# where
# -n (--num_rep) k: number of simulation repetitions
# -d (--debug)  m: debug verbose level
#                  (0 = none; 1 = perf measures; 2 = sample paths)
# -t (--trial): trial runs (uses different random numbers than production runs)
# -p (--p_type) a: parameter to determines performance measure
#   a=0 - expected time to first fill-up of all beds
#   a=1 - expected fraction of time in [0,30] with at least 7 idle beds
#   a=2 - probability that time to first fill-up is at most 55 days
#   a=3 - expected frac. of arrivals in [0,30] resulting in a hospital transfer
#
# NOTES:
#     Beds are numbered starting from 0 (because of Python's 0-based indexing)
#     Event order in clock-reading vector:
#         c[0]: arrival of patient
#         c[1]: departure from SCU bed 0
#         ...
#         c[NUM_BEDS_SCU]: departure from SCU bed NUM_BEDS_SCU - 1
#         c[NUM_BEDS_SCU+1]: departure from ICU bed 0
#         ...
#         c[NUM_BEDS_SCU+NUM_BEDS_ICU]: departure from ICU bed NUM_BEDS_ICU - 1
#         c[NUM_BEDS_SCU+NUM_BEDS_ICU+1]: critical event for SCU bed 0
#         ...
#         c[NUM_BEDS_SCU+NUM_BEDS_ICU+NUM_BEDS_SCU]: critical event for
#             SCU bed NUM_BEDS_SCU - 1
#
# Uses the rngStream module (the module should be in the python path or
# the same folder)


class Estimator:
    """ Computes point estimates and confidence intervals """
    def __init__(self, z, conf_str):
        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 = float(z)  # quantile for normal Confidence Interval
        self.conf_str = conf_str  # string of form "xx%" for xx% conf. int.

    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, epsilon, relative=True):
        var = self.get_variance()
        if var == 0:
            return "UNAVAILABLE (Need at least 2 pilot reps)"
        width = self.get_mean() * epsilon if relative else epsilon
        return int(1+(var * self.z * self.z)/(width * width))


def print_sc(time, mm, nn, c, istar, model_params, arr_scu):
    """ Format state, clock reading vector, and trigger event for printing """
    cf = "{:5.2f}"  # format string for clock readings
    m = " ".join(map(str, mm))
    n = " ".join(map(str, nn))
    s_str = "({0} | {1})".format(m,n)
    nb_scu = model_params["num_beds_scu"]
    nb_icu = model_params["num_beds_icu"]
    c0 = cf.format(c[0])
    c1 = " ".join([cf.format(x) for x in c[1:(nb_scu + 1)]])
    c2 = " ".join([cf.format(x) for x in 
                  c[(nb_scu + 1):(nb_scu + nb_icu + 1)]])
    c3 = " ".join([cf.format(x)for x in
                  c[(nb_scu + nb_icu + 1):(2 * nb_scu + nb_icu + 1)]])
    c_str = "({0} | {1} | {2} | {3})".format(c0, c1, c2, c3)
    end_str = ""
    if istar >= 0:
        end_str = "i* = {:d}".format(istar)
    if istar == 0:
        end_str = end_str + (" (SCU)" if arr_scu else " (ICU)")
    return "T = {0:8.4f}: {1}\n{2} {3}\n\n".format(time, s_str, c_str, end_str)


def do_rep(unigens, est, model_params, p_type):
    """ Executes one replication of the emergency room simulation """
    arr_pat, scu_dep, icu_dep = unigens  # unpack random number generators
    nb_scu = model_params["num_beds_scu"]
    nb_icu = model_params["num_beds_icu"]

    # initialize time, state, clocks, and performance measure
    time = 0   # simulation time
    perf_meas = 0
    mm = [0] * nb_scu
    nn = [0] * nb_icu
    total_beds = nb_scu + nb_icu
    num_arrivals = 0
    c_length = 1 + 2 * nb_scu + nb_icu
    c = [-1] * c_length  # empty clock reading vector
    # schedule first arrival
    u = arr_pat[0].RandU01()
    c[0] = model_params["arr_param"] * sqrt(u)

    done = False  # True if simulation replication is completed

    while not done:  # simulate the GSMP
        # print("<"+str(time)+">")

        # compute holding time in current state and trigger event
        clist = [(index, value) for index, value in enumerate(c) if value >= 0]
        istar, htime = min(clist, key=lambda p: p[1])

        # set (s',c') = (s, c) as initial default
        mmp = list(mm)
        nnp = list(nn)
        cp = list(c)

        # compute udated clock reading vector (will update it more below)
        cp = [x - htime if x > 0 else x for x in c]
        cp[istar] = -1

        # compute next state s' and next clocks c' (i.e., schedule events)
        # branch according to which event triggers next transition
        if istar == 0:  # arrival of patient
            num_arrivals += 1
            u = arr_pat[1].RandU01()  # determine type of arriving patient
            if u <= model_params["scu_prob"]:  # arriving patient goes to SCU
                arr_scu = True
                # find empty SCU bed, if any:
                i = next((k for k, x in enumerate(mm) if x == 0), None)
                if i is not None:  # there is an empty SCU bed
                    mmp[i] = 1  # mark bed as occupied
                    # schedule departure
                    u1 = scu_dep[i][0].RandU01()
                    u2 = scu_dep[i][1].RandU01()
                    cp[i+1] = -(log(u1)+log(u2))/model_params["stay_rate"]
                    # schedule critical event
                    u1 = scu_crit[i][0].RandU01()
                    u2 = scu_crit[i][1].RandU01()
                    cp[nb_scu + nb_icu + 1 + i] = 0.5 * (u1 + u2) + 0.5
            else:  # arriving patient goes to ICU
                arr_scu = False
                # find empty ICU bed, if any:
                i = next((k for k, x in enumerate(nn) if x == 0), None)
                if i is not None:  # there is an empty ICU bed
                    nnp[i] = 1  # mark bed as occupied
                    # schedule ICU departure
                    u1 = icu_dep[i][0].RandU01()
                    u2 = icu_dep[i][1].RandU01()
                    cp[nb_scu+1+i] = \
                        -(log(u1)+log(u2))/model_params["stay_rate"]
            u = arr_pat[0].RandU01()  # schedule next arrival
            cp[0] = model_params["arr_param"] * sqrt(u)
        elif istar <= nb_scu:  # an SCU patient just departed
            j = istar - 1  # compute bed number (0-indexed) of dep. patient
            mmp[j] = 0  # mark bed as empty
            cp[nb_scu+nb_icu+1+j] = -1  # cancel critical event
        elif istar <= nb_scu + nb_icu:  # an ICU patient just departed
            j = istar - nb_scu - 1  # compute bed # (0-indexed) of dep. patient
            nnp[j] = 0  # mark bed as empty
        else:  # there has been a critical event in the SCU
            j = istar - nb_scu - nb_icu - 1  # bed # (0-indexed) of crit. ptnt
            mmp[j] = 0  # mark critical SCU bed as empty
            cp[j+1] = -1  # cancel SCU departure
            # find empty ICU bed, if any:
            i = next((k for k, x in enumerate(nn) if x == 0), None)
            if i is not None:  # there is an empty ICU bed
                nnp[i] = 1  # mark bed as occupied
                u1 = icu_dep[i][0].RandU01()
                u2 = icu_dep[i][1].RandU01()
                cp[nb_scu+1+i] = \
                    -(log(u1)+log(u2))/model_params["stay_rate"]
        # end of (s', c') processing

        #  debug print of (s, c, e*)
        verbose_print(2, print_sc(time, mm, nn, c, istar, model_params,
                                  arr_scu))

        # compute/update performance measure and check for termination
        if p_type == 0:  # expected time to first fill-up of all beds
            if sum(mmp + nnp) == total_beds:  # we're done
                perf_meas = time + htime
                done = True
        elif p_type == 1:  # expected frac. time in [0,30] with >= 7 idle beds
            if time + htime > model_params["sim_time"]:  # we're done
                htime = model_params["sim_time"] - time
                done = True
            if total_beds - sum(mm+nn) >= 7:
                perf_meas += htime
        elif p_type == 2:  # prob. that time to first fill-up is <= 55 days
            if sum(mmp + nnp) == total_beds:  # we're done
                perf_meas = 1 if time + htime <= 55.0 else 0
                done = True
        elif p_type == 3:  # E[frac. arrivals in [0, 30] that get transferred]
            if time + htime > model_params["sim_time"]:  # we're done
                done = True
            if (not done) and istar == 0 \
                    and sum(mm + nn) == sum(mmp + nnp):
                perf_meas += 1  # transfer has occurred

        # update state, clocks, and time
        mm = list(mmp)  # make sure to actually copy the data
        nn = list(nnp)
        c = list(cp)
        time += htime
        # end of GSMP simulation ("while not done" loop)

    # finalize performance measure computation if necessary
    if p_type == 1:
        perf_meas /= model_params["sim_time"]
    if p_type == 3:
        perf_meas = perf_meas / num_arrivals if num_arrivals > 0 else 0

    # final detailed debug print
    verbose_print(2, print_sc(time, mm, nn, c, -1, model_params, False))

    # update statistics and return
    est.process_next_val(perf_meas)
    return perf_meas

if __name__ == "__main__":
    # parse command line arguments
    parser = ArgumentParser(
        description="hospitalg -num_rep k [-debug m -trial -p_type a]")
    parser.add_argument('-n', '--num_rep', 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="Select performance measure",
                        default=0)

    sysargs = parser.parse_args()

    est = Estimator(2.576, "99%")  # 95% CI

    epsilon = 0.01  # Determines the width of the CI

    relative = True  # Determines relative vs absolute accuracy

    # ensure that results are reproducible
    seed0 = [12345, 12345, 12345, 12345, 12345, 12345]
    rngStream.SetPackageSeed(seed0)

    model_params = {
        "num_beds_scu": 5,
        "num_beds_icu": 4,
        "arr_param": 0.08,  # parameter a of interarrival distn F(x) = (x/a)^2
        "scu_prob": 0.68,   # prob. that an arriving patient goes to SCU
        "stay_rate": 8.0,   # rate for erlang(2) distn (SCU stay and ICU stay
        "sim_time": 30      # time horizon of interest
        }

    perf_name = (
        "E[time to first fill-up]",
        "E[frac time that >= 7 idle]",
        "P(time to first fill-up <= 55)",
        "E[fraction of transfers]"
        )

    # Instantiate the random number generators & package them up
    arr_pat = []
    for i in range(2):
        arr_pat.append(rngStream.RngStream("arr{:d}".format(i)))
        if not sysargs.trial:
            arr_pat[i].ResetNextSubstream()
    scu_dep = []
    scu_crit = []
    for i in range(model_params["num_beds_scu"]):
        scu_dep.append([])
        scu_crit.append([])
        for j in range(2):
            scu_dep[i].append(
                rngStream.RngStream("SCUdep{0:d}{0:d}".format(i,j)))
            scu_crit[i].append(
                rngStream.RngStream("SCUcrit{0:d}{0:d}".format(i,j)))
            if not sysargs.trial:
                scu_dep[i][j].ResetNextSubstream()
                scu_crit[i][j].ResetNextSubstream()
    icu_dep = []
    for i in range(model_params["num_beds_icu"]):
        icu_dep.append([])
        for j in range(2):
            icu_dep[i].append(
                rngStream.RngStream("ICUdep{0:d}{0:d}".format(i,j)))
            if not sysargs.trial:
                icu_dep[i][j].ResetNextSubstream()
    unigens = [arr_pat, scu_dep, icu_dep]

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

    # run simulation repetitions and collect stats
    print("\n")
    for rep in range(int(sysargs.num_rep)):
        val = do_rep(unigens, est, model_params, int(sysargs.p_type))
        verbose_print(1, "Repetition {0:3d} : {1:.4f}\n\n".format(rep+1, val))

    # print results
    verbose_print(0, "{}:\n".format(perf_name[int(sysargs.p_type)]))
    verbose_print(0, " Point Estimate: {:.2f}\n".format(est.get_mean()))
    verbose_print(0, " {}".format(est.get_conf_interval()))
    print("\n")
    if sysargs.trial:
        if relative:
            print("Est. # of reps for +/-", 100 * epsilon,
                  "percent accuracy: ", est.get_num_trials(epsilon, relative))
        else:
            print("Est. # of reps for +/-", epsilon, "accuracy: ",
                  est.get_num_trials(epsilon, relative))