import os

from .diagnostic import VerosDiagnostic
from .. import veros_class_method
from ..variables import Variable

ENERGY_VARIABLES = dict(
    # mean energy content
    k_m=Variable("Mean kinetic energy", [], "J", "Mean kinetic energy",
                 output=True, write_to_restart=True),
    Hd_m=Variable("Mean dynamic enthalpy", [], "J", "Mean dynamic enthalpy",
                  output=True, write_to_restart=True),
    eke_m=Variable("Meso-scale eddy energy", [], "J", "Meso-scale eddy energy",
                   output=True, write_to_restart=True),
    iw_m=Variable("Internal wave energy", [], "J", "Internal wave energy",
                  output=True, write_to_restart=True),
    tke_m=Variable("Turbulent kinetic energy", [], "J", "Turbulent kinetic energy",
                   output=True, write_to_restart=True),

    # energy changes
    dE_tot_m=Variable("Change of total energy", [], "W", "Change of total energy",
                      output=True, write_to_restart=True),
    dk_m=Variable("Change of KE", [], "W", "Change of kinetic energy",
                  output=True, write_to_restart=True),
    dHd_m=Variable("Change of Hd", [], "W", "Change of dynamic enthalpy",
                   output=True, write_to_restart=True),
    deke_m=Variable("Change of EKE", [], "W", "Change of meso-scale eddy energy",
                    output=True, write_to_restart=True),
    diw_m=Variable("Change of E_iw", [], "W", "Change of internal wave energy",
                   output=True, write_to_restart=True),
    dtke_m=Variable("Change of TKE", [], "W", "Change of tubulent kinetic energy",
                    output=True, write_to_restart=True),

    # dissipation
    ke_diss_m=Variable("Dissipation of KE", [], "W", "Dissipation of kinetic energy",
                       output=True, write_to_restart=True),
    Hd_diss_m=Variable("Dissipation of Hd", [], "W", "Dissipation of dynamic enthalpy",
                       output=True, write_to_restart=True),
    eke_diss_m=Variable("Dissipation of EKE", [], "W", "Dissipation of meso-scale eddy energy",
                        output=True, write_to_restart=True),
    iw_diss_m=Variable("Dissipation of E_iw", [], "W", "Dissipation of internal wave energy",
                       output=True, write_to_restart=True),
    tke_diss_m=Variable("Dissipation of TKE", [], "W", "Dissipation of turbulent kinetic energy",
                        output=True, write_to_restart=True),
    adv_diss_m=Variable("Dissipation by advection", [], "W", "Dissipation by advection",
                        output=True, write_to_restart=True),

    # external forcing
    wind_m=Variable("Wind work", [], "W", "Wind work",
                    output=True, write_to_restart=True),
    dHd_sources_m=Variable("Hd production by ext. sources", [], "W",
                           "Dynamic enthalpy production through external sources",
                           output=True, write_to_restart=True),
    iw_forc_m=Variable("External forcing of E_iw", [], "W",
                       "External forcing of internal wave energy",
                       output=True, write_to_restart=True),
    tke_forc_m=Variable("External forcing of TKE", [], "W",
                        "External forcing of turbulent kinetic energy",
                        output=True, write_to_restart=True),

    # exchange
    ke_hd_m=Variable("Exchange KE -> Hd", [], "W",
                     "Exchange between kinetic energy and dynamic enthalpy",
                     output=True, write_to_restart=True),
    ke_tke_m=Variable("Exchange KE -> TKE by vert. friction", [], "W",
                      "Exchange between kinetic energy and turbulent kinetic energy by vertical friction",
                      output=True, write_to_restart=True),
    ke_iw_m=Variable("Exchange KE -> IW by bottom friction", [], "W",
                     "Exchange between kinetic energy and internal wave energy by bottom friction",
                     output=True, write_to_restart=True),
    tke_hd_m=Variable("Exchange TKE -> Hd by vertical mixing", [], "W",
                      "Exchange between turbulent kinetic energy and dynamic enthalpy by vertical mixing",
                      output=True, write_to_restart=True),
    ke_eke_m=Variable("Exchange KE -> EKE by lateral friction", [], "W",
                      "Exchange between kinetic energy and eddy kinetic energy by lateral friction",
                      output=True, write_to_restart=True),
    hd_eke_m=Variable("Exchange Hd -> EKE by GM and lateral mixing", [], "W",
                      "Exchange between dynamic enthalpy and eddy kinetic energy by GM and lateral mixing",
                      output=True, write_to_restart=True),
    eke_tke_m=Variable("Exchange EKE -> TKE", [], "W",
                       "Exchange between eddy and turbulent kinetic energy",
                       output=True, write_to_restart=True),
    eke_iw_m=Variable("Exchange EKE -> IW", [], "W",
                      "Exchange between eddy kinetic energy and internal wave energy",
                      output=True, write_to_restart=True),

    # cabbeling
    cabb_m=Variable("Cabbeling by vertical mixing", [], "W",
                    "Cabbeling by vertical mixing",
                    output=True, write_to_restart=True),
    cabb_iso_m=Variable("Cabbeling by isopycnal mixing", [], "W",
                        "Cabbeling by isopycnal mixing",
                        output=True, write_to_restart=True),
)


class Energy(VerosDiagnostic):
    """Diagnose globally averaged energy cycle. Also averages energy in time.
    """
    name = "energy" #:
    output_path = "{identifier}.energy.nc"  #: File to write to. May contain format strings that are replaced with Veros attributes.
    output_frequency = None  #: Frequency (in seconds) in which output is written.
    sampling_frequency = None  #: Frequency (in seconds) in which variables are accumulated.
    variables = ENERGY_VARIABLES

    @veros_class_method
    def initialize(self, vs):
        self.nitts = 0
        for var in self.variables.keys():
            setattr(self, var, 0.)

        output_variables = {key: val for key, val in self.variables.items() if val.output}
        self.initialize_output(vs, output_variables)

    @veros_class_method
    def diagnose(self, vs):
        # changes of dynamic enthalpy
        vol_t = vs.area_t[2:-2, 2:-2, np.newaxis] \
            * vs.dzt[np.newaxis, np.newaxis, :] \
            * vs.maskT[2:-2, 2:-2, :]
        dP_iso = np.sum(vol_t * vs.grav / vs.rho_0
                        * (-vs.int_drhodT[2:-2, 2:-2, :, vs.tau]
                           * vs.dtemp_iso[2:-2, 2:-2, :]
                           - vs.int_drhodS[2:-2, 2:-2, :, vs.tau]
                           * vs.dsalt_iso[2:-2, 2:-2, :]))
        dP_hmix = np.sum(vol_t * vs.grav / vs.rho_0
                         * (-vs.int_drhodT[2:-2, 2:-2, :, vs.tau]
                            * vs.dtemp_hmix[2:-2, 2:-2, :]
                            - vs.int_drhodS[2:-2, 2:-2, :, vs.tau]
                            * vs.dsalt_hmix[2:-2, 2:-2, :]))
        dP_vmix = np.sum(vol_t * vs.grav / vs.rho_0
                         * (-vs.int_drhodT[2:-2, 2:-2, :, vs.tau]
                            * vs.dtemp_vmix[2:-2, 2:-2, :]
                            - vs.int_drhodS[2:-2, 2:-2, :, vs.tau]
                            * vs.dsalt_vmix[2:-2, 2:-2, :]))
        dP_m = np.sum(vol_t * vs.grav / vs.rho_0
                      * (-vs.int_drhodT[2:-2, 2:-2, :, vs.tau]
                          * vs.dtemp[2:-2, 2:-2, :, vs.tau]
                          - vs.int_drhodS[2:-2, 2:-2, :, vs.tau]
                          * vs.dsalt[2:-2, 2:-2, :, vs.tau]))
        dP_m_all = dP_m + dP_vmix + dP_hmix + dP_iso

        # changes of kinetic energy
        vol_u = vs.area_u[2:-2, 2:-2, np.newaxis] \
            * vs.dzt[np.newaxis, np.newaxis, :]
        vol_v = vs.area_v[2:-2, 2:-2, np.newaxis] \
            * vs.dzt[np.newaxis, np.newaxis, :]
        k_m = np.sum(vol_t * 0.5 * (0.5 * (vs.u[2:-2, 2:-2, :, vs.tau] ** 2
                                           + vs.u[1:-3, 2:-2, :, vs.tau] ** 2)
                                    + 0.5 * (vs.v[2:-2, 2:-2, :, vs.tau] ** 2)
                                    + vs.v[2:-2, 1:-3, :, vs.tau] ** 2))
        p_m = np.sum(vol_t * vs.Hd[2:-2, 2:-2, :, vs.tau])
        dk_m = np.sum(vs.u[2:-2, 2:-2, :, vs.tau] * vs.du[2:-2, 2:-2, :, vs.tau] * vol_u
                      + vs.v[2:-2, 2:-2, :, vs.tau]
                      * vs.dv[2:-2, 2:-2, :, vs.tau] * vol_v
                      + vs.u[2:-2, 2:-2, :, vs.tau] * vs.du_mix[2:-2, 2:-2, :] * vol_u
                      + vs.v[2:-2, 2:-2, :, vs.tau] * vs.dv_mix[2:-2, 2:-2, :] * vol_v)

        corm = np.sum(vs.u[2:-2, 2:-2, :, vs.tau] * vs.du_cor[2:-2, 2:-2, :] * vol_u
                      + vs.v[2:-2, 2:-2, :, vs.tau] * vs.dv_cor[2:-2, 2:-2, :] * vol_v)
        k_e_adv = np.sum(vs.u[2:-2, 2:-2, :, vs.tau] * vs.du_adv[2:-2, 2:-2, :]
                         * vol_u * vs.maskU[2:-2, 2:-2, :]
                         + vs.v[2:-2, 2:-2, :, vs.tau] * vs.dv_adv[2:-2, 2:-2, :]
                         * vol_v * vs.maskV[2:-2, 2:-2, :])

        # K*Nsqr and KE and dyn. enthalpy dissipation
        vol_w = vs.area_t[2:-2, 2:-2, np.newaxis] * vs.dzw[np.newaxis, np.newaxis, :] \
            * vs.maskW[2:-2, 2:-2, :]
        vol_w[:, :, -1] *= 0.5

        def mean_w(var):
            return np.sum(var[2:-2, 2:-2, :] * vol_w)

        mdiss_vmix = mean_w(vs.P_diss_v)
        mdiss_nonlin = mean_w(vs.P_diss_nonlin)
        mdiss_adv = mean_w(vs.P_diss_adv)
        mdiss_hmix = mean_w(vs.P_diss_hmix)
        mdiss_comp = mean_w(vs.P_diss_comp)
        mdiss_iso = mean_w(vs.P_diss_iso)
        mdiss_skew = mean_w(vs.P_diss_skew)
        mdiss_sources = mean_w(vs.P_diss_sources)

        mdiss_h = mean_w(vs.K_diss_h)
        mdiss_v = mean_w(vs.K_diss_v)
        mdiss_gm = mean_w(vs.K_diss_gm)
        mdiss_bot = mean_w(vs.K_diss_bot)

        wrhom = np.sum(-vs.area_t[2:-2, 2:-2, np.newaxis] * vs.maskW[2:-2, 2:-2, :-1]
                       * (vs.p_hydro[2:-2, 2:-2, 1:] - vs.p_hydro[2:-2, 2:-2, :-1])
                       * vs.w[2:-2, 2:-2, :-1, vs.tau])

        # wind work
        wind = np.sum(vs.u[2:-2, 2:-2, -1, vs.tau] * vs.surface_taux[2:-2, 2:-2]
                      * vs.maskU[2:-2, 2:-2, -1] * vs.area_u[2:-2, 2:-2]
                      + vs.v[2:-2, 2:-2, -1, vs.tau] * vs.surface_tauy[2:-2, 2:-2]
                      * vs.maskV[2:-2, 2:-2, -1] * vs.area_v[2:-2, 2:-2])

        # meso-scale energy
        if vs.enable_eke:
            eke_m = mean_w(vs.eke[..., vs.tau])
            deke_m = np.sum(vol_w * (vs.eke[2:-2, 2:-2, :, vs.taup1]
                                     - vs.eke[2:-2, 2:-2, :, vs.tau])
                            / vs.dt_tracer)
            eke_diss = mean_w(vs.eke_diss_iw)
            eke_diss_tke = mean_w(vs.eke_diss_tke)
        else:
            eke_m = deke_m = eke_diss_tke = 0.
            eke_diss = mdiss_gm + mdiss_h + mdiss_skew
            if not vs.enable_store_cabbeling_heat:
                eke_diss += -mdiss_hmix - mdiss_iso

        # small-scale energy
        if vs.enable_tke:
            tke_m = mean_w(vs.tke[..., vs.tau])
            dtke_m = mean_w((vs.tke[..., vs.taup1]
                             - vs.tke[..., vs.tau])
                            / vs.dt_tke)
            tke_diss = mean_w(vs.tke_diss)
            tke_forc = np.sum(vs.area_t[2:-2, 2:-2] * vs.maskW[2:-2, 2:-2, -1]
                              * (vs.forc_tke_surface[2:-2, 2:-2] + vs.tke_surf_corr[2:-2, 2:-2]))
        else:
            tke_m = dtke_m = tke_diss = tke_forc = 0.

        # internal wave energy
        if vs.enable_idemix:
            iw_m = mean_w(vs.E_iw[..., vs.tau])
            diw_m = np.sum(vol_w * (vs.E_iw[2:-2, 2:-2, :, vs.taup1]
                                    - vs.E_iw[2:-2, 2:-2, :, vs.tau])
                           / vs.dt_tracer)
            iw_diss = mean_w(vs.iw_diss)

            k = np.maximum(1, vs.kbot[2:-2, 2:-2]) - 1
            mask = k[:, :, np.newaxis] == np.arange(vs.nz)[np.newaxis, np.newaxis, :]
            iwforc = np.sum(vs.area_t[2:-2, 2:-2]
                            * (vs.forc_iw_surface[2:-2, 2:-2] * vs.maskW[2:-2, 2:-2, -1]
                               + np.sum(mask * vs.forc_iw_bottom[2:-2, 2:-2, np.newaxis]
                                        * vs.maskW[2:-2, 2:-2, :], axis=2)))
        else:
            iw_m = diw_m = iwforc = 0.
            iw_diss = eke_diss

        # store results
        self.k_m += k_m
        self.Hd_m += p_m
        self.eke_m += eke_m
        self.iw_m += iw_m
        self.tke_m += tke_m

        self.dk_m += dk_m
        self.dHd_m += dP_m_all + mdiss_sources
        self.deke_m += deke_m
        self.diw_m += diw_m
        self.dtke_m += dtke_m
        self.dE_tot_m += self.dk_m + self.dHd_m + self.deke_m + self.diw_m + self.dtke_m

        self.wind_m += wind
        self.dHd_sources_m += mdiss_sources
        self.iw_forc_m += iwforc
        self.tke_forc_m += tke_forc

        self.ke_diss_m += mdiss_h + mdiss_v + mdiss_gm + mdiss_bot
        self.Hd_diss_m += mdiss_vmix + mdiss_nonlin + mdiss_hmix + mdiss_adv \
            + mdiss_iso + mdiss_skew
        self.eke_diss_m += eke_diss + eke_diss_tke
        self.iw_diss_m += iw_diss
        self.tke_diss_m += tke_diss
        self.adv_diss_m += mdiss_adv

        self.ke_hd_m += wrhom
        self.ke_eke_m += mdiss_h + mdiss_gm
        self.hd_eke_m += -mdiss_skew
        self.ke_tke_m += mdiss_v
        self.tke_hd_m += -mdiss_vmix - mdiss_adv
        if vs.enable_store_bottom_friction_tke:
            self.ke_tke_m += mdiss_bot
        else:
            self.ke_iw_m += mdiss_bot
        self.eke_tke_m += eke_diss_tke
        self.eke_iw_m += eke_diss
        if not vs.enable_store_cabbeling_heat:
            self.hd_eke_m += -mdiss_hmix - mdiss_iso
            self.tke_hd_m += -mdiss_nonlin

        self.cabb_m += mdiss_nonlin
        self.cabb_iso_m += mdiss_hmix + mdiss_iso

        self.nitts += 1

    @veros_class_method
    def output(self, vs):
        self.nitts = float(self.nitts or 1)
        output_variables = {key: val for key, val in self.variables.items() if val.output}
        output_data = {key: getattr(self, key) * vs.rho_0 /
                       self.nitts for key in output_variables.keys()}
        if not os.path.isfile(self.get_output_file_name(vs)):
            self.initialize_output(vs, output_variables)
        self.write_output(vs, output_variables, output_data)

        for key in output_variables.keys():
            setattr(self, key, 0.)
        self.nitts = 0

    @veros_class_method
    def read_restart(self, vs):
        attributes, variables = self.read_h5_restart(vs)
        if attributes:
            for key, val in attributes.items():
                setattr(self, key, val)

    @veros_class_method
    def write_restart(self, vs, outfile):
        restart_data = {key: getattr(self, key)
                        for key, val in self.variables.items() if val.write_to_restart}
        restart_data.update({"nitts": self.nitts})
        self.write_h5_restart(vs, restart_data, {}, {}, outfile)
