# Policy Evaluation Class for pc-gym
import numpy as np
import matplotlib.pyplot as plt
from .oracle import oracle


class policy_eval:
    """
    Policy Evaluation Class

    Inputs: Environment, policy and number of policy repitions

    Outputs: Plots of states/control/constraints/setpoints (complete),
             return distribution (incomplete), expected return (incomplete),
             oracle trajectories (incomplete) and lower confidence bounds (incomplete)
    """

    def __init__(
        self,
        make_env,
        policies,
        reps,
        env_params,
        oracle=False,
        MPC_params=False,
        cons_viol=False,
    ):
        self.make_env = make_env
        self.env_params = env_params
        self.env = make_env(env_params)
        self.policies = policies
        self.n_pi = len(policies)
        self.reps = reps
        self.oracle = oracle
        self.cons_viol = cons_viol

        self.MPC_params = MPC_params

    def rollout(self, policy_i):
        """
        Rollout the policy for N steps and return the total reward, states and actions

        Input:
            policy - policy to be rolled out

        Outputs:
            total_reward - total reward obtained
            states - states obtained from rollout
            actions - actions obtained from rollout

        """

        total_reward = 0
        s_rollout = np.zeros((self.env.Nx, self.env.N))
        actions = np.zeros((self.env.env_params["a_space"]["low"].shape[0], self.env.N))

        o, r = self.env.reset()
        total_reward = r["r_init"]
        s_rollout[:, 0] = (o + 1) * (
            self.env.observation_space_base.high - self.env.observation_space_base.low
        ) / 2 + self.env.observation_space_base.low
        for i in range(self.env.N - 1):
            a, _s = policy_i.predict(
                o, deterministic=True
            )  # Rollout with a deterministic policy
            o, r, term, trunc, info = self.env.step(a)

            actions[:, i] = (a + 1) * (
                self.env.env_params["a_space"]["high"]
                - self.env.env_params["a_space"]["low"]
            ) / 2 + self.env.env_params["a_space"]["low"]
            s_rollout[:, i + 1] = (o + 1) * (
                self.env.observation_space_base.high - self.env.observation_space_base.low
            ) / 2 + self.env.observation_space_base.low

            total_reward += r

        if self.env.constraint_active:
            cons_info = info["cons_info"]
        else:
            cons_info = np.zeros((1, self.env.N, 1))
        a, _s = policy_i.predict(o, deterministic=True)
        actions[:, self.env.N - 1] = (a + 1) * (
            self.env.env_params["a_space"]["high"]
            - self.env.env_params["a_space"]["low"]
        ) / 2 + self.env.env_params["a_space"]["low"]

        return total_reward, s_rollout, actions, cons_info

    def get_rollouts(self):
        """
        Function to plot the rollout of the policy

        Inputs:
            policy - policy to be rolled out
            reps - number of rollouts to be performed

        Outputs:
            Plot of states and actions with setpoints and constraints if they exist]

        """
        data = {}
        action_space_shape = self.env.env_params["a_space"]["low"].shape[0]
        num_states = self.env.Nx

        # Collect Oracle data
        if self.oracle:
            r_opt = np.zeros((1, self.reps))
            x_opt = np.zeros((self.env.Nx_oracle, self.env.N, self.reps))
            u_opt = np.zeros((self.env.Nu, self.env.N, self.reps))
            oracle_instance = oracle(self.make_env, self.env_params, self.MPC_params)
            for i in range(self.reps):
                x_opt[:, :, i], u_opt[:, :, i] = oracle_instance.mpc()
                for k in self.env.SP:
                    state_i = self.env.model.info()["states"].index(k)
                    r_scale = self.env_params.get("r_scale", {})
                    r_opt[:, i] += (
                        np.sum((x_opt[state_i, :, i] - self.env.SP[k]) ** 2)
                        * -1
                        * r_scale.get(k, 1)
                    )
            data.update({"oracle": {"r": r_opt, "x": x_opt, "u": u_opt}})

        # Collect RL rollouts for all policies
        for pi_name, pi_i in self.policies.items():
            states = np.zeros((num_states, self.env.N, self.reps))
            actions = np.zeros((action_space_shape, self.env.N, self.reps))
            rew = np.zeros((1, self.reps))
            try:
                cons_info = np.zeros((self.env.n_con, self.env.N, 1, self.reps))
            except Exception:
                cons_info = np.zeros((1, self.env.N, 1, self.reps))
            for r_i in range(self.reps):
                (
                    rew[:, r_i],
                    states[:, :, r_i],
                    actions[:, :, r_i],
                    cons_info[:, :, :, r_i],
                ) = self.rollout(pi_i)
            data.update({pi_name: {"r": rew, "x": states, "u": actions}})
            if self.env.constraint_active:
                data[pi_name].update({"g": cons_info})
        self.data = data
        return data

    def plot_data(self, data, reward_dist=False):
        t = np.linspace(0, self.env.tsim, self.env.N)
        len_d = 0

        if self.env.disturbance_active:
            len_d = len(self.env.model.info()["disturbances"])

        col = ["tab:red", "tab:purple", "tab:olive", "tab:gray", "tab:cyan"]
        if self.n_pi > len(col):
            raise ValueError(
                f"Number of policies ({self.n_pi}) is greater than the number of available colors ({len(col)})"
            )

        plt.figure(figsize=(10, 2 * (self.env.Nx_oracle + self.env.Nu - self.env.Nd)))
        for i in range(self.env.Nx_oracle):
            plt.subplot(self.env.Nx_oracle + self.env.Nu - self.env.Nd, 1, i + 1)
            for ind, (pi_name, pi_i) in enumerate(self.policies.items()):
                plt.plot(
                    t,
                    np.median(data[pi_name]["x"][i, :, :], axis=1),
                    color=col[ind],
                    lw=3,
                    label=self.env.model.info()["states"][i] + " (" + pi_name + ")",
                )
                plt.gca().fill_between(
                    t,
                    np.min(data[pi_name]["x"][i, :, :], axis=1),
                    np.max(data[pi_name]["x"][i, :, :], axis=1),
                    color=col[ind],
                    alpha=0.2,
                    edgecolor="none",
                )
            if self.oracle:
                plt.plot(
                    t,
                    np.median(data["oracle"]["x"][i, :, :], axis=1),
                    color="tab:blue",
                    lw=3,
                    label="Oracle " + self.env.model.info()["states"][i],
                )
                plt.gca().fill_between(
                    t,
                    np.min(data["oracle"]["x"][i, :, :], axis=1),
                    np.max(data["oracle"]["x"][i, :, :], axis=1),
                    color="tab:blue",
                    alpha=0.2,
                    edgecolor="none",
                )
            if self.env.model.info()["states"][i] in self.env.SP:
                plt.step(
                    t,
                    self.env.SP[self.env.model.info()["states"][i]],
                    where="post",
                    color="black",
                    linestyle="--",
                    label="Set Point",
                )
            if self.env.constraint_active:
                if self.env.model.info()["states"][i] in self.env.constraints:
                    plt.hlines(
                        self.env.constraints[self.env.model.info()["states"][i]],
                        0,
                        self.env.tsim,
                        color="black",
                        label="Constraint",
                    )
            plt.ylabel(self.env.model.info()["states"][i])
            plt.xlabel("Time (min)")
            plt.legend(loc="best")
            plt.grid("True")
            plt.xlim(min(t), max(t))

        for j in range(self.env.Nu - len_d):
            plt.subplot(
                self.env.Nx_oracle + self.env.Nu - self.env.Nd,
                1,
                j + self.env.Nx_oracle + 1,
            )
            for ind, (pi_name, pi_i) in enumerate(self.policies.items()):
                plt.step(
                    t,
                    np.median(data[pi_name]["u"][j, :, :], axis=1),
                    color=col[ind],
                    lw=3,
                    label=self.env.model.info()["inputs"][j] + " (" + pi_name + ")",
                )
            if self.oracle:
                plt.step(
                    t,
                    np.median(data["oracle"]["u"][j, :, :], axis=1),
                    color="tab:blue",
                    lw=3,
                    label="Oracle " + str(self.env.model.info()["inputs"][j]),
                )
            if self.env.constraint_active:
                for con_i in self.env.constraints:
                    if self.env.model.info()["inputs"][j] == con_i:
                        plt.hlines(
                            self.env.constraints[self.env.model.info()["inputs"][j]],
                            0,
                            self.env.tsim,
                            "black",
                            label="Constraint",
                        )
            plt.ylabel(self.env.model.info()["inputs"][j])
            plt.xlabel("Time (min)")
            plt.legend(loc="best")
            plt.grid("True")
            plt.xlim(min(t), max(t))

        if self.env.disturbance_active:
            for k in self.env.disturbances.keys():
                i = 1
                if self.env.disturbances[k].any() is not None:
                    plt.subplot(
                        self.env.Nx_oracle + self.env.Nu - self.env.Nd,
                        1,
                        i + j + self.env.Nx_oracle + 1,
                    )
                    plt.step(t, self.env.disturbances[k], color="tab:orange", label=k)
                    plt.xlabel("Time (min)")
                    plt.ylabel(k)
                    plt.xlim(min(t), max(t))
                    i += 1
        plt.tight_layout()
        plt.show()

        if self.cons_viol:
            plt.figure(figsize=(12, 3 * self.env.n_con))
            con_i = 0
            for i, con in enumerate(self.env.constraints):
                for j in range(len(self.env.constraints[str(con)])):
                    plt.subplot(self.env.n_con, 1, con_i + 1)
                    plt.title(f"{con} Constraint")
                    for ind, (pi_name, pi_i) in enumerate(self.policies.items()):
                        plt.step(
                            t,
                            np.sum(data[pi_name]["g"][con_i, :, :, :], axis=2),
                            color=col[ind],
                            label=f"{con} ({pi_name}) Violation (Sum over Repetitions)",
                        )
                    plt.grid("True")
                    plt.xlabel("Time (min)")
                    plt.ylabel(con)
                    plt.xlim(min(t), max(t))
                    plt.legend(loc="best")
                    con_i += 1
            plt.tight_layout()
            plt.show()

        if reward_dist:
            plt.figure(figsize=(12, 8))
            plt.grid(True, linestyle="--", alpha=0.6)
            all_data = np.concatenate([data[key]["r"].flatten() for key in data.keys()])

            min_value = np.min(all_data)
            max_value = np.max(all_data)

            bins = np.linspace(min_value, max_value, self.reps)
            if self.oracle:
                plt.hist(
                    data["oracle"]["r"].flatten(),
                    bins=bins,
                    color="tab:blue",
                    alpha=0.5,
                    label="Oracle",
                    edgecolor="black",
                )
            for ind, (pi_name, pi_i) in enumerate(self.policies.items()):
                plt.hist(
                    data[pi_name]["r"].flatten(),
                    bins=bins,
                    color=col[ind],
                    alpha=0.5,
                    label=pi_name,
                    edgecolor="black",
                )

            plt.xlabel("Return", fontsize=14)
            plt.ylabel("Frequency", fontsize=14)
            plt.title("Distribution of Expected Return", fontsize=16)
            plt.legend(fontsize=12)

            plt.show()

        return
