import numpy as np


class ODEStates:
    def __init__(self, t) -> None:

        self.t = t

        n_points = len(self.t)

        # start filling up from 0th index
        self.index = 0

        # state values
        self.S = np.zeros(n_points)

        self.ERR = np.zeros(n_points)
        self.ERS = np.zeros(n_points)
        self.ESR = np.zeros(n_points)
        self.ESS = np.zeros(n_points)

        self.IRR = np.zeros(n_points)
        self.IRS = np.zeros(n_points)
        self.ISR = np.zeros(n_points)
        self.ISS = np.zeros(n_points)

        self.R = np.zeros(n_points)

        self.PRR = np.zeros(n_points)
        self.PRS = np.zeros(n_points)
        self.PSR = np.zeros(n_points)
        self.PSS = np.zeros(n_points)

        self.fung_1 = np.zeros(n_points)
        self.fung_2 = np.zeros(n_points)

    def update_y(self, y):

        ind = self.index

        length = y.shape[1]

        self.S[ind:ind+length] = y[0, :]

        self.ERR[ind:ind+length] = y[1, :]
        self.ERS[ind:ind+length] = y[2, :]
        self.ESR[ind:ind+length] = y[3, :]
        self.ESS[ind:ind+length] = y[4, :]

        self.IRR[ind:ind+length] = y[5, :]
        self.IRS[ind:ind+length] = y[6, :]
        self.ISR[ind:ind+length] = y[7, :]
        self.ISS[ind:ind+length] = y[8, :]

        self.R[ind:ind+length] = y[9, :]

        self.PRR[ind:ind+length] = y[10, :]
        self.PRS[ind:ind+length] = y[11, :]
        self.PSR[ind:ind+length] = y[12, :]
        self.PSS[ind:ind+length] = y[13, :]

        self.fung_1[ind:ind+length] = y[14, :]
        self.fung_2[ind:ind+length] = y[15, :]

        self.index = ind + length

    def delete_unnecessary_vars(self):
        delattr(self, "index")


class SimOutput:
    def __init__(self, t_vec) -> None:
        self.states = ODEStates(t_vec)

        self.final_res_vec_dict = None
        self.end_freqs = None
        self.yield_val = None


class SingleTacticOutput:
    """Output for RunSingleTactic

    class with following attributes:
        - failure_year : year yield below 95

        - yield_vec : array of yields for each year

        - res_vec_dict : dict with keys f1, f2 - total resistance to fcides

        - start_freqs : dict with keys RR, RS, SR, SS - frequencies, 
        post-sex from previous year, ready for start of season

        - failure_year : dict with keys RR, RS, SR, SS - frequencies, 
        end of season, pre-sex

        - states_list : contains list for each year, each element
        containing values of variables over time (e.g. S, E_RR, etc at
        each time point)
    """

    def __init__(self,
                 yield_thresh,
                 res_props_in,
                 strain_names,
                 n_years,
                 df_yield) -> None:

        self.yield_thresh = yield_thresh
        self.n_years = n_years
        self.df_yield = df_yield
        self.strain_names = strain_names

        self.failure_year = 0

        self.yield_vec = np.zeros(n_years)

        self.res_vec_dict = self._init_res_vec_dict(res_props_in)

        # post-sex from previous year, ready for start of season
        self.start_freqs = self._init_freq_dict()

        # end of season, pre-sex
        self.end_freqs = self._init_freq_dict()

        self.states_list = []

    def _init_freq_dict(self):
        out = {}

        for key in self.strain_names:
            out[key] = np.zeros(self.n_years+1)

        return out

    def _init_res_vec_dict(self, res_props):
        out = {}
        keys = ['f1', 'f2']

        for key in keys:
            out[key] = np.zeros(self.n_years+1)
            # set first year
            out[key][0] = res_props[key]

        return out

    def add_new_sim_output(self, sim_out, yr):

        self.yield_vec[yr] = 100*(sim_out.yield_val/self.df_yield)

        self._update_end_freqs(sim_out, yr)

        self._update_res_vec_dict(sim_out, yr+1)

        self._update_failure_year(yr)

        self.states_list.append(sim_out.states)

    def _update_res_vec_dict(self, sim_out, yr):
        for key in ['f1', 'f2']:
            self.res_vec_dict[key][yr] = sim_out.final_res_vec_dict[key]

    def update_start_freqs(self, values, yr):
        for key in self.strain_names:
            self.start_freqs[key][yr] = values[key]

    def _update_end_freqs(self, sim_out, yr):
        for key in self.end_freqs.keys():
            self.end_freqs[key][yr] = sim_out.end_freqs[key]

    def _update_failure_year(self, yr):
        """
        Set failure year if:
        - yield is below threshold
        - is first time it has dropped below threshold
        """

        if ((self.yield_vec[yr] < self.yield_thresh) and
                (self.failure_year == 0)):
            self.failure_year = yr+1

    def delete_unnecessary_vars(self):
        delattr(self, "yield_thresh")
        delattr(self, "n_years")
        delattr(self, "df_yield")
        delattr(self, "strain_names")


class GridTacticOutput:
    def __init__(self, n_doses, n_years) -> None:

        self.LTY = np.zeros((n_doses, n_doses))
        self.TY = np.zeros((n_doses, n_doses))
        self.FY = np.zeros((n_doses, n_doses))

        self.yield_array = np.zeros((n_doses, n_doses, n_years))

        fung_keys = ['f1', 'f2']
        self.res_vec_DA = self._get_dict_of_zero_arrays(
            fung_keys, (n_doses, n_doses, n_years+1))

        strain_keys = ['RR', 'RS', 'SR', 'SS']
        self.start_freqs_DA = self._get_dict_of_zero_arrays(
            strain_keys, (n_doses, n_doses, n_years+1))
        self.end_freqs_DA = self._get_dict_of_zero_arrays(
            strain_keys, (n_doses, n_doses, n_years+1))

    @staticmethod
    def _get_dict_of_zero_arrays(keys, shape):
        out = {}
        for key in keys:
            out[key] = np.zeros(shape)
        return out

    def update_dicts_of_arrays(self, data, f1_ind, f2_ind):
        mydata = vars(data)

        self.res_vec_DA = self._update_dict_array_this_dose(
            self.res_vec_DA, mydata["res_vec_dict"],
            f1_ind, f2_ind)

        self.start_freqs_DA = self._update_dict_array_this_dose(
            self.start_freqs_DA, mydata["start_freqs"],
            f1_ind, f2_ind)

        self.end_freqs_DA = self._update_dict_array_this_dose(
            self.end_freqs_DA, mydata["end_freqs"],
            f1_ind, f2_ind)

    def _update_dict_array_this_dose(self, to_update, calculated, f1_ind, f2_ind):

        for key_ in to_update.keys():
            to_update[key_][f1_ind, f2_ind, :] = calculated[key_]

        return to_update

import numpy as np


class ODEStates:
    def __init__(self, t) -> None:

        self.t = t

        n_points = len(self.t)

        # start filling up from 0th index
        self.index = 0

        # state values
        self.S = np.zeros(n_points)

        self.ERR = np.zeros(n_points)
        self.ERS = np.zeros(n_points)
        self.ESR = np.zeros(n_points)
        self.ESS = np.zeros(n_points)

        self.IRR = np.zeros(n_points)
        self.IRS = np.zeros(n_points)
        self.ISR = np.zeros(n_points)
        self.ISS = np.zeros(n_points)

        self.R = np.zeros(n_points)

        self.PRR = np.zeros(n_points)
        self.PRS = np.zeros(n_points)
        self.PSR = np.zeros(n_points)
        self.PSS = np.zeros(n_points)

        self.fung_1 = np.zeros(n_points)
        self.fung_2 = np.zeros(n_points)

    def update_y(self, y):

        ind = self.index

        length = y.shape[1]

        self.S[ind:ind+length] = y[0, :]

        self.ERR[ind:ind+length] = y[1, :]
        self.ERS[ind:ind+length] = y[2, :]
        self.ESR[ind:ind+length] = y[3, :]
        self.ESS[ind:ind+length] = y[4, :]

        self.IRR[ind:ind+length] = y[5, :]
        self.IRS[ind:ind+length] = y[6, :]
        self.ISR[ind:ind+length] = y[7, :]
        self.ISS[ind:ind+length] = y[8, :]

        self.R[ind:ind+length] = y[9, :]

        self.PRR[ind:ind+length] = y[10, :]
        self.PRS[ind:ind+length] = y[11, :]
        self.PSR[ind:ind+length] = y[12, :]
        self.PSS[ind:ind+length] = y[13, :]

        self.fung_1[ind:ind+length] = y[14, :]
        self.fung_2[ind:ind+length] = y[15, :]

        self.index = ind + length

    def delete_unnecessary_vars(self):
        delattr(self, "index")


class SimOutput:
    def __init__(self, t_vec) -> None:
        self.states = ODEStates(t_vec)

        self.final_res_vec_dict = None
        self.end_freqs = None
        self.yield_val = None


class SingleTacticOutput:
    """Output for RunSingleTactic

    class with following attributes:
        - failure_year : year yield below 95

        - yield_vec : array of yields for each year

        - res_vec_dict : dict with keys f1, f2 - total resistance to fcides

        - start_freqs : dict with keys RR, RS, SR, SS - frequencies, 
        post-sex from previous year, ready for start of season

        - failure_year : dict with keys RR, RS, SR, SS - frequencies, 
        end of season, pre-sex

        - states_list : contains list for each year, each element
        containing values of variables over time (e.g. S, E_RR, etc at
        each time point)
    """

    def __init__(self,
                 yield_thresh,
                 res_props_in,
                 strain_names,
                 n_years,
                 df_yield) -> None:

        self.yield_thresh = yield_thresh
        self.n_years = n_years
        self.df_yield = df_yield
        self.strain_names = strain_names

        self.failure_year = 0

        self.yield_vec = np.zeros(n_years)

        self.res_vec_dict = self._init_res_vec_dict(res_props_in)

        # post-sex from previous year, ready for start of season
        self.start_freqs = self._init_freq_dict()

        # end of season, pre-sex
        self.end_freqs = self._init_freq_dict()

        self.states_list = []

    def _init_freq_dict(self):
        out = {}

        for key in self.strain_names:
            out[key] = np.zeros(self.n_years+1)

        return out

    def _init_res_vec_dict(self, res_props):
        out = {}
        keys = ['f1', 'f2']

        for key in keys:
            out[key] = np.zeros(self.n_years+1)
            # set first year
            out[key][0] = res_props[key]

        return out

    def add_new_sim_output(self, sim_out, yr):

        self.yield_vec[yr] = 100*(sim_out.yield_val/self.df_yield)

        self._update_end_freqs(sim_out, yr)

        self._update_res_vec_dict(sim_out, yr+1)

        self._update_failure_year(yr)

        self.states_list.append(sim_out.states)

    def _update_res_vec_dict(self, sim_out, yr):
        for key in ['f1', 'f2']:
            self.res_vec_dict[key][yr] = sim_out.final_res_vec_dict[key]

    def update_start_freqs(self, values, yr):
        for key in self.strain_names:
            self.start_freqs[key][yr] = values[key]

    def _update_end_freqs(self, sim_out, yr):
        for key in self.end_freqs.keys():
            self.end_freqs[key][yr] = sim_out.end_freqs[key]

    def _update_failure_year(self, yr):
        """
        Set failure year if:
        - yield is below threshold
        - is first time it has dropped below threshold
        """

        if ((self.yield_vec[yr] < self.yield_thresh) and
                (self.failure_year == 0)):
            self.failure_year = yr+1

    def delete_unnecessary_vars(self):
        delattr(self, "yield_thresh")
        delattr(self, "n_years")
        delattr(self, "df_yield")
        delattr(self, "strain_names")


class GridTacticOutput:
    def __init__(self, n_doses, n_years) -> None:

        self.LTY = np.zeros((n_doses, n_doses))
        self.TY = np.zeros((n_doses, n_doses))
        self.FY = np.zeros((n_doses, n_doses))

        self.yield_array = np.zeros((n_doses, n_doses, n_years))

        fung_keys = ['f1', 'f2']
        self.res_vec_DA = self._get_dict_of_zero_arrays(
            fung_keys, (n_doses, n_doses, n_years+1))

        strain_keys = ['RR', 'RS', 'SR', 'SS']
        self.start_freqs_DA = self._get_dict_of_zero_arrays(
            strain_keys, (n_doses, n_doses, n_years+1))
        self.end_freqs_DA = self._get_dict_of_zero_arrays(
            strain_keys, (n_doses, n_doses, n_years+1))

    @staticmethod
    def _get_dict_of_zero_arrays(keys, shape):
        out = {}
        for key in keys:
            out[key] = np.zeros(shape)
        return out

    def update_dicts_of_arrays(self, data, f1_ind, f2_ind):
        mydata = vars(data)

        self.res_vec_DA = self._update_dict_array_this_dose(
            self.res_vec_DA, mydata["res_vec_dict"],
            f1_ind, f2_ind)

        self.start_freqs_DA = self._update_dict_array_this_dose(
            self.start_freqs_DA, mydata["start_freqs"],
            f1_ind, f2_ind)

        self.end_freqs_DA = self._update_dict_array_this_dose(
            self.end_freqs_DA, mydata["end_freqs"],
            f1_ind, f2_ind)

    def _update_dict_array_this_dose(self, to_update, calculated, f1_ind, f2_ind):

        for key_ in to_update.keys():
            to_update[key_][f1_ind, f2_ind, :] = calculated[key_]

        return to_update

# start filling up from 0th index # state values Output for RunSingleTactic      class with following attributes:         - failure_year : year yield below 95          - yield_vec : array of yields for each year          - res_vec_dict : dict with keys f1, f2 - total resistance to fcides          - start_freqs : dict with keys RR, RS, SR, SS - frequencies,          post-sex from previous year, ready for start of season          - failure_year : dict with keys RR, RS, SR, SS - frequencies,          end of season, pre-sex          - states_list : contains list for each year, each element         containing values of variables over time (e.g. S, E_RR, etc at         each time point) # post-sex from previous year, ready for start of season # end of season, pre-sex # set first year Set failure year if:         - yield is below threshold         - is first time it has dropped below threshold