# PyMC3
# SimPY?

"""
Based on BAC, goal is to use data to generate a good simulation of the queue allowing us to determine expected wait times under different scenarios.

Data is generated by a combination of arriving flights. Arrival times are sampled from a beta distribution (varying by flight?).
Processing times (per lane) are sum of exponential (no check or with check).

Model is built with the same structure using PyMC3 to estimate parameters.

Finally simulations are run with SimPY and animated under different conditions.

Verify with goodness of fit test + metrics?
"""


# Model system
# Tasks : Sim data
# Plot sim data
# PyMC3 recover sim parameters
# Create sim with SimPy
# Plot animated simulation
# Add settings

# Does Kolmovorov-Smirnov-Test work to compare sim to real world?



import simpy  # https://simpy.readthedocs.io/en/latest/simpy_intro/index.html
from scipy.stats import gamma
import matplotlib.pyplot as plt
from typing import List, Tuple
import pandas as pd
import streamlit as st
import altair as alt
import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pymc3 as pm
import pandas as pd

# Monitor:  Passenger - arrival_time, wait_time, processing_time
#           Queue - Queue length per time step


        


#-------------------------------------------Simulation----------------------------------------------------
class Passenger:

    def __init__(self, env, station, name: str, result_store: List, walk_dist, check_dist):
        self.env = env
        self.station = station
        self.name = name
        self.arrival_time = walk_dist.rvs()
        self.processing_time = check_dist.rvs()
        # Start the run process everytime an instance is created.
        self.action = env.process(self.run())
        self.data = {"arrival_time": self.arrival_time, "wait_time": -1, "processing_time": self.processing_time}
        self.result_store = result_store

    def run(self):
        # print(f"{self.name}: Walking to passport control at {self.env.now}")
        yield self.env.timeout(self.arrival_time)

        queue_arrival_time = self.env.now
        # print(f"{self.name}: Arriving at passport control at {self.env.now}")
        with self.station.request() as req:
            yield req
            wait_time = self.env.now - queue_arrival_time
            self.data['wait_time'] = wait_time

            # print(f"{self.name}: Checking Passport at {self.env.now}")
            yield self.env.timeout(self.processing_time)

            # print(f"{self.name}: Leaving passport control at {self.env.now}")

        self.dump(self.data)

    def dump(self, data: dict):
        self.result_store.add_result(data)


class Plotter:

    def __init__(self, env, result_store, plot_container, freq=2):
        self.env = env
        self.plot_container = plot_container
        self.result_store = result_store
        self.action = env.process(self.run())
        self.freq = freq
        print("Creating Plotter")

    def run(self):
        print("Running Plotter")
        while True:
            try:
                data = self.result_store.get_results()[1]
                chart = DataDisplay().plot_queue_length(data)
                self.plot_container.altair_chart(chart.properties(
                        width=20 * data['time'].max()
                    ).encode(
                        y=alt.Y('queue_length',
                            scale=alt.Scale(domain=(0, data['queue_length'].max()))
                            )
                        )
                )
                print(f"Plotting - {self.env.now}")
            except Exception as e:
                print(e)
            yield self.env.timeout(self.freq)




class ResultStore:

    def __init__(self):
        self.results = []
        self.queue_lengths = []

    def add_result(self, data):
        self.results.append(data)

    def add_queue_data(self, data):
        self.queue_lengths.append(data)

    def get_results(self):
        return pd.DataFrame(self.results), pd.DataFrame(self.queue_lengths)

    def save(self):
        pd.DataFrame(self.results).to_csv("passenger_data.csv")
        pd.DataFrame(self.queue_lengths).to_csv("queue_data.csv")


class MonitoredResource(simpy.Resource):
    def __init__(self, result_store, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.result_store = result_store

    def request(self, *args, **kwargs):
        self.result_store.add_queue_data({"time": self._env.now, "queue_length": len(self.queue)})
        return super().request(*args, **kwargs)

    def release(self, *args, **kwargs):
        self.result_store.add_queue_data({"time": self._env.now, "queue_length": len(self.queue)})
        return super().release(*args, **kwargs)


class OnePlaneSimulation:

    def __init__(self, n_passengers: int, n_lanes: int, walk_a: float, check_a: float, walk_scale: float, check_scale: float, max_steps=1000, plot_container=None):
        self.n_passengers = n_passengers
        self.n_lanes = n_lanes
        self.walk_dist = gamma(walk_a, scale=walk_scale)
        self.check_dist = gamma(check_a, scale=check_scale)
        self.result_store = ResultStore()
        self.max_steps = max_steps
        self.plot_container = plot_container

    def run(self):
        timeout=500
        if self.plot_container is not None:
            env = simpy.rt.RealtimeEnvironment(factor=0.5)
            timeout=50
        else:
            env = simpy.Environment()

        station = MonitoredResource(env=env, result_store=self.result_store, capacity=self.n_lanes)
        names = [str(i) for i in range(self.n_passengers)]
        [Passenger(env, station, name, self.result_store, self.walk_dist, self.check_dist) for name in names]  # init passengers

        if self.plot_container is not None:
            print("New Plotter")
            Plotter(env, self.result_store, self.plot_container)

        env.run(until=env.timeout(timeout))
        # self.result_store.save()
        return self.get_results()

    def get_results(self):
        return self.result_store.get_results()

    def save_results(self):
        self.result_store.save()


#-------------------------------------------Modelling Fitting----------------------------------------------------


class Modeller:

    def __init__(self):
        print(f"Running on PyMC3 v{pm.__version__}")

    def model_arrival_time(self, arrival_times):
        arrival_model = pm.Model()

        with arrival_model:

            # Priors for unknown model parameters
            alpha = pm.HalfCauchy("alpha", beta=5)
            scale = pm.HalfCauchy("scale", beta=5)
            beta = 1 / scale

            # Likelihood (sampling distribution) of observations
            Y_obs = pm.Gamma("Y_obs", alpha=alpha, beta=beta, observed=arrival_times)

            # draw 2000 posterior samples
            trace = pm.sample(2000, return_inferencedata=False, cores=1)
            plot = az.plot_trace(trace)
            summary = az.summary(trace, round_to=2)
            print("Finished Modelling Arrivals")

        return {"alpha": summary.loc["alpha", "mean"], "scale": summary.loc["scale", "mean"], "summary": summary, "plot": plot}

    def model_processing_time(self, processing_times):
        arrival_model = pm.Model()

        with arrival_model:

            # Priors for unknown model parameters
            alpha = pm.HalfNormal("alpha", sigma=100)
            scale = pm.HalfNormal("scale", sigma=100)
            beta = 1 / scale

            # Likelihood (sampling distribution) of observations
            Y_obs = pm.Gamma("Y_obs", alpha=alpha, beta=beta, observed=processing_times)

            # draw 2000 posterior samples
            trace = pm.sample(2000, return_inferencedata=False, cores=1)
            plot = az.plot_trace(trace)
            summary = az.summary(trace, round_to=2)
            print("Finished Processing Arrivals")

        return {"alpha": summary.loc["alpha", "mean"], "scale": summary.loc["scale", "mean"], "summary": summary, "plot": plot}

#-------------------------------------------Validation----------------------------------------------------


#-------------------------------------------Displays----------------------------------------------------

class SummaryDisplay:

    def show(self, data):
        passenger_data, queue_data = data
        
        mean_queue_time = passenger_data["wait_time"].mean()
        max_queue_time = passenger_data["wait_time"].max()
        median_queue_time = passenger_data["wait_time"].median()
        mean_procesing_time = passenger_data["processing_time"].mean()

        st.markdown(f"""
        ## KPIs
        Mean queue time: {mean_queue_time:.0f}\t
        Median queue time: {median_queue_time:.0f}\t
        Max queue time: {max_queue_time:.0f}\t

        ## Other Info:
        Mean Processing Time: {mean_procesing_time:.1f}
        """)

class DataDisplay:

    def show(self, data):
        chart = self.get_chart(data)
        st.altair_chart(chart, use_container_width=True)

    def get_chart(self, data):
        passenger_data, queue_data = data
        arr_chart = self.plot_arrival_times(passenger_data)
        wait_chart = self.plot_wait_times(passenger_data)
        check_chart = self.plot_processing_time(passenger_data)
        queue_chart = self.plot_queue_length(queue_data)
        return (queue_chart & wait_chart) | (arr_chart & check_chart)

    def plot_wait_times(self, df):
        chart = alt.Chart(df).mark_bar().encode(
            alt.X("wait_time:Q", bin=alt.Bin(maxbins=30)),
            y='count()',
        ).properties(
            title='Wait Times',
            width=200
        )
        return chart

    def plot_arrival_times(self, df):
        chart = alt.Chart(df).mark_line().encode(
            alt.X("arrival_time:Q", bin=alt.Bin(maxbins=100)),
            y='count()',
        ).properties(
            title='Arrival Times',
            width=200
        )
        return chart

    def plot_processing_time(self, df):
        chart = alt.Chart(df).mark_bar().encode(
            alt.X("processing_time:Q", bin=alt.Bin(maxbins=30)),
            y='count()',
        ).properties(
            title='Check Durations',
            width=200
        )
        return chart

    def plot_queue_length(self, df):
        chart = alt.Chart(df).mark_line().encode(
            x='time',
            y='queue_length',
        ).properties(
            title='Queue Length',
            width=200
        )
        return chart

    def combine_charts(self, charts, columns):
        return alt.vconcat(alt.hconcat(chart for chart in charts[i * columns: (i + 1) * columns]) for i in range(len(charts) // columns) )


class FitDisplay:

    def __init__(self):
        self.arrival_results = None
        self.processing_results = None

    def show(self):
        fit = st.button('Fit')

        if fit:
            self.fit()

            arrival_fig = plt.figure(1,2)
            if self.arrival_results:
                st.markdown(f"Arrivals estimates: shape: { self.arrival_results['alpha']}, scale: { self.arrival_results['scale']}")
                [arrival_fig.add_subplot(p) for p in self.arrival_results['plot'].flatten()]
            st.pyplot(arrival_fig)

            processing_fig = plt.figure(1, 2)
            if self.processing_results:
                st.markdown(f"Processing estimates: shape: {self.processing_results['alpha']}, scale: {self.processing_results['scale']}")
                # [processing_fig.add_subplot(p) for p in self.processing_results['plot'].flatten()]
            # st.pyplot(processing_fig)

    def fit(self):
        data = st.session_state.screen.get_data()
        if data is None:
            return None

        print(data)
        self.arrival_results = Modeller().model_arrival_time(data[0]['arrival_time'])
        print(self.arrival_results)
        self.processing_results = Modeller().model_processing_time(data[0]['processing_time'])







class Settings:

    def show(self):
        st.button('Rerun')
        self.settings = {}
        self.settings['n_passengers'] = st.number_input("Number of Passengers", min_value=1, format="%i",value=100)
        self.settings['n_lanes'] = st.number_input("Number of Lanes", min_value=1, format="%i", value=5)
        self.settings['animate'] = st.checkbox("Animate Result")
        
        with st.expander("Detailed Settings"):
            self.settings['walk_a']  = st.number_input("Walk shape: Controls how spread out arrivals are", min_value=1, format="%i", value=5)
            self.settings['walk_scale'] = st.number_input("Walk Scale: Controls how long it takes for passengers to arrive", min_value=0, format="%i", value=1)
            self.settings['check_a'] = st.number_input("Check shape: Controls how spread out processing times are", min_value=1, format="%i", value=1)
            self.settings['check_scale'] = st.number_input("Check Scale: Controls how long checks take.", min_value=0, format="%i", value=1)

    def get_settings(self) -> Tuple[int, int, float, float]:
        return self.settings

#-------------------------------------------Screen----------------------------------------------------]


class QueueScreen:

    def __init__(self):
        self.data = None
        self.settings = Settings()
        
    def show(self):
        st.markdown(
        """
        # Scenario: Queue Simulation
        """
        )
        col_settings, col_plots, col_fit = st.columns(3)
        with col_settings:
            self.settings.show()

        with col_plots:
            plot_loc = st.empty()

        settings = self.settings.get_settings()
        settings['plot_container'] = plot_loc if settings['animate'] else None
        del settings['animate']
        sim = OnePlaneSimulation(**settings)
        self.data = sim.run()

        with col_plots:
            plot_loc.empty()
            SummaryDisplay().show(self.data)
            DataDisplay().show(self.data)

        with col_fit:
            FitDisplay().show()

    def get_data(self):
        return self.data

def init_state():
    if "screen" not in st.session_state:
        st.session_state.screen = QueueScreen()


if __name__ == "__main__":
    st.set_page_config(layout="wide")


    init_state()
    st.session_state.screen.show()