# coding=utf-8
import numpy as np

from earthdiagnostics.box import Box
from earthdiagnostics.constants import Basins
from earthdiagnostics.diagnostic import Diagnostic, DiagnosticIntOption, DiagnosticBasinOption
from earthdiagnostics.modelingrealm import ModelingRealms
from earthdiagnostics.utils import Utils, TempFile


class HeatContentLayer(Diagnostic):
    """
    Point-wise Ocean Heat Content in a specified ocean thickness (J/m-2)

    :original author: Isabel Andreu Burillo
    :contributor: Virginie Guemas <virginie.guemas@bsc.es>
    :contributor: Eleftheria Exarchou <eleftheria.exarchou@bsc.es>
    :contributor: Javier Vegas-Regidor<javier.vegas@bsc.es>

    :created: June 2012
    :last modified: June 2016

    :param data_manager: data management object
    :type data_manager: DataManager
    :param startdate: startdate
    :type startdate: str
    :param member: member number
    :type member: int
    :param chunk: chunk's number
    :type chunk: int
    :param box: box to use for the calculations
    :type box: Box
    """

    alias = 'ohclayer'
    "Diagnostic alias for the configuration file"

    def __init__(self, data_manager, startdate, member, chunk, box, weight, min_level, max_level):
        Diagnostic.__init__(self, data_manager)
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.box = box
        self.weight = weight
        self.min_level = min_level
        self.max_level = max_level
        self.required_vars = ['so', 'mlotst']
        self.generated_vars = ['scvertsum']

    def __str__(self):
        return 'Heat content layer Startdate: {0} Member: {1} Chunk: {2} Box: {3}'.format(self.startdate, self.member,
                                                                                          self.chunk, self.box)

    @classmethod
    def generate_jobs(cls, diags, options):
        """
        Create a job for each chunk to compute the diagnostic

        :param diags: Diagnostics manager class
        :type diags: Diags
        :param options: minimum depth, maximum depth, basin=Global
        :type options: list[str]
        """
        options_available = (DiagnosticIntOption('min_depth'),
                             DiagnosticIntOption('max_depth'),
                             DiagnosticBasinOption('basin', Basins().Global))
        options = cls.process_options(options, options_available)

        box = Box(True)
        box.min_depth = options['min_depth']
        box.max_depth = options['max_depth']
        job_list = list()

        handler = Utils.open_cdf('mesh_zgr.nc')
        # mask = Utils.get_mask(options['basin'])
        mask = handler.variables['tmask'][:]

        if 'e3t' in handler.variables:
            e3t = handler.variables['e3t'][:]
        elif 'e3t_0' in handler.variables:
            e3t = handler.variables['e3t_0'][:]
        else:
            raise Exception('e3t variable can not be found')

        if 'gdepw' in handler.variables:
            depth = handler.variables['gdepw'][:]
        elif 'gdepw_0' in handler.variables:
            depth = handler.variables['gdepw_0'][:]
        else:
            raise Exception('gdepw variable can not be found')

        e3t_3d = e3t.shape != depth.shape
        if e3t_3d:
            mask = e3t_3d * mask
        else:
            e3t = e3t[0, :]

        while len(depth.shape) < 4:
            depth = np.expand_dims(depth, -1)
        handler.close()

        def calculate_weight(e3t_point, depth_point, mask_point):
            """Calculates the weight for each cell"""
            if not mask_point:
                return 0
            top = depth_point
            bottom = top + e3t_point
            if bottom < box.min_depth or top > box.max_depth:
                return 0
            else:
                if top < box.min_depth:
                    top = box.min_depth
                if bottom > box.max_depth:
                    bottom = box.max_depth

                return (bottom - top) * 1020 * 4000

        calc = np.vectorize(calculate_weight, otypes='f')
        weight = calc(e3t, depth, mask)

        # Now we will reduce to the levels with any weight != 0 to avoid loading too much data on memory
        levels = weight.shape[1]
        min_level = 0
        while min_level < levels and not weight[:, min_level, :].any():
            min_level += 1
        max_level = min_level
        while max_level < (levels - 1) and weight[:, max_level + 1, :].any():
            max_level += 1
        weight = weight[:, min_level:max_level, :]

        for startdate, member, chunk in diags.config.experiment.get_chunk_list():
            job_list.append(HeatContentLayer(diags.data_manager, startdate, member, chunk, box,
                                             weight, min_level, max_level))
        return job_list

    def request_data(self):
        """Request data required by the diagnostic"""
        self.thetao = self.request_chunk(ModelingRealms.ocean, 'thetao', self.startdate, self.member, self.chunk)

    def declare_data_generated(self):
        """Declare data to be generated by the diagnostic"""
        self.heatc = self.declare_chunk(ModelingRealms.ocean, 'heatc', self.startdate, self.member, self.chunk,
                                        box=self.box)

    def compute(self):
        """Run the diagnostic"""
        nco = Utils.nco
        thetao_file = TempFile.get()
        results = TempFile.get()

        Utils.copy_file(self.thetao.local_file, thetao_file)

        handler = Utils.open_cdf(thetao_file)
        Utils.convert_units(handler.variables['thetao'], 'K')
        heatc_sl = np.sum(handler.variables['thetao'][:, self.min_level:self.max_level, :] * self.weight, 1)
        handler.sync()
        handler.renameVariable('thetao', 'heatc_sl')
        handler.close()

        nco.ncks(input=thetao_file, output=results, options=('-O -v lon,lat,time',))
        Utils.rename_variables(results, {'x': 'i', 'y': 'j'}, False, True)
        handler_results = Utils.open_cdf(results)
        handler_results.createVariable('heatc', float, ('time', 'j', 'i'), fill_value=1.e20)
        handler_results.sync()
        handler_results.variables['heatc'][:] = heatc_sl
        handler_results.close()

        Utils.setminmax(results, 'heatc')
        self.heatc.set_local_file(results)