from importlib import reload

import lab5_funs

reload(lab5_funs)
from lab5_funs import Integrator
from collections import namedtuple
import numpy as np


class Integ54(Integrator):
    def set_yinit(self):
        #
        # read in 'albedo_white chi S0 L albedo_black R albedo_ground'
        #
        uservars = namedtuple("uservars", self.config["uservars"].keys())
        self.uservars = uservars(**self.config["uservars"])
        #
        # read in 'whiteconc blackconc'
        #
        initvars = namedtuple("initvars", self.config["initvars"].keys())
        self.initvars = initvars(**self.config["initvars"])
        self.yinit = np.array([self.initvars.whiteconc, self.initvars.blackconc])
        self.nvars = len(self.yinit)
        return None

    def __init__(self, coeff_file_name):
        super().__init__(coeff_file_name)
        self.set_yinit()

    def find_temp(self, yvals):
        """
            Calculate the temperatures over the white and black daisies
            and the planetary equilibrium temperature given the daisy fractions
            
            input:  yvals -- array of dimension [2] with the white [0] and black [1]
                    daisy fractiion
            output:  white temperature (K), black temperature (K), equilibrium temperature (K)
        """
        sigma = 5.67e-8  # Stefan Boltzman constant W/m^2/K^4
        user = self.uservars
        bare = 1.0 - yvals[0] - yvals[1]
        albedo_p = (
            bare * user.albedo_ground
            + yvals[0] * user.albedo_white
            + yvals[1] * user.albedo_black
        )
        Te_4 = user.S0 / 4.0 * user.L * (1.0 - albedo_p) / sigma
        temp_e = Te_4 ** 0.25
        eta = user.R * user.L * user.S0 / (4.0 * sigma)
        temp_b = (eta * (albedo_p - user.albedo_black) + Te_4) ** 0.25
        temp_w = (eta * (albedo_p - user.albedo_white) + Te_4) ** 0.25
        return (temp_w, temp_b, temp_e)

    def derivs5(self, y, t):
        """y[0]=fraction white daisies
           y[1]=fraction black daisies
           no feedback between daisies and
           albedo_p (set to ground albedo)
        """
        temp_w, temp_b, temp_e = self.find_temp(y)

        if temp_b >= 277.5 and temp_b <= 312.5:
            beta_b = 1.0 - 0.003265 * (295.0 - temp_b) ** 2.0
        else:
            beta_b = 0.0

        if temp_w >= 277.5 and temp_w <= 312.5:
            beta_w = 1.0 - 0.003265 * (295.0 - temp_w) ** 2.0
        else:
            beta_w = 0.0
        user = self.uservars
        bare = 1.0 - y[0] - y[1]
        # create a 1 x 2 element vector to hold the derivitive
        f = np.empty_like(y)
        f[0] = y[0] * (beta_w * bare - user.chi)
        f[1] = y[1] * (beta_b * bare - user.chi)
        return f


if __name__ == "__main__":
    import matplotlib.pyplot as plt
    import pandas as pd

    theSolver = Integ54("init_files/initial.yaml")
    timevals, yvals, errorlist = theSolver.timeloop5fixed()
    daisies = pd.DataFrame(yvals, columns=["white", "black"])

    thefig, theAx = plt.subplots(1, 1)
    line1, = theAx.plot(timevals, daisies["white"])
    line2, = theAx.plot(timevals, daisies["black"])
    line1.set(linestyle="--", color="r", label="white")
    line2.set(linestyle="--", color="k", label="black")
    theAx.set_title("lab 5 interactive 4, initial conditions")
    theAx.set_xlabel("time")
    theAx.set_ylabel("fractional coverage")
    out = theAx.legend(loc="center right")
    thefig.savefig("example54.png")
    plt.show()
