import numpy as np
from earthdiagnostics.basins import Basins
from earthdiagnostics.diagnostic import Diagnostic
from earthdiagnostics.box import Box
from earthdiagnostics.utils import Utils, TempFile


class AreaMoc(Diagnostic):
    """
    Compute an Atlantic MOC index by averaging the meridional overturning
    in a latitude band between 1km and 2km
    or any other index averaging the meridional overturning in
    a given basin and a given domain

    :original author: Virginie Guemas <virginie.guemas@bsc.es>
    :contributor: Javier Vegas-Regidor<javier.vegas@bsc.es>

    :created: March 2012
    :last modified: June 2016

    """

    def __init__(self, data_manager, startdate, member, chunk, basin, box):
        Diagnostic.__init__(self, data_manager)
        self.basin = basin
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.required_vars = ['vo']
        self.generated_vars = ['vsftmyz']
        self.box = box

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

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

        :param diags: Diagnostics manager class
        :type diags: Diags
        :param options: minimum latitude, maximum latitude, minimum depth, maximum depth, basin=Global
        :type options: tuple
        :return:
        """
        num_options = len(options) - 1
        if num_options < 4:
            raise Exception('You must specify the box to use')
        if num_options > 5:
            raise Exception('You must specify between 4 and 5 parameters for area moc diagnostic')
        box = Box()
        box.min_lat = int(options[1])
        box.max_lat = int(options[2])
        box.min_depth = int(options[3])
        box.max_depth = int(options[4])
        if num_options > 4:
            basin = Basins.parse(options[5])
        else:
            basin = Basins.Global

        job_list = list()
        for startdate in diags.startdates:
            for member in diags.members:
                for chunk in range(1, diags.chunks + 1):
                    job_list.append(AreaMoc(diags.data_manager, startdate, member, chunk, basin, box))
        return job_list

    def compute(self):
        nco = Utils.nco
        cdo = Utils.cdo
        temp2 = TempFile.get()

        temp = self.data_manager.get_file('ocean', 'vsftmyz', self.startdate, self.member, self.chunk)

        handler = Utils.openCdf(temp)
        if 'i' in handler.dimensions:
            handler.close()
            nco.ncwa(input=temp, output=temp, options='-O -a i')
            handler = Utils.openCdf(temp)

        basin_index = np.where(handler.variables['basin'][:] == self.basin.fullname)
        lat_values = handler.variables['lat'][:]
        lat_type = handler.variables['lat'].dtype
        lat_units = handler.variables['lat'].units
        lat_long_name = handler.variables['lat'].long_name

        handler.close()

        if len(basin_index) == 0:
            raise Exception('Basin {0} not defined in file')
        basin_index = basin_index[0][0]
        nco.ncks(input=temp, output=temp, options='-O -d basin,{0}'.format(basin_index))
        # To remove basin dimension
        nco.ncwa(input=temp, output=temp, options='-O -a basin')

        nco.ncks(input=temp, output=temp, options='-O -v vsftmyz,time,lev')

        handler = Utils.openCdf(temp)
        handler.renameDimension('j', 'lat')
        lat_variable = handler.createVariable('lat', lat_type, 'lat')
        lat_variable[:] = lat_values[:]
        lat_variable.units = lat_units
        lat_variable.long_name = lat_long_name

        handler.close()

        nco.ncks(input=temp, output=temp2,
                 options='-O -d lev,{0:.1f},{1:.1f} -d lat,{2:.1f},{3:.1f}'.format(self.box.min_depth,
                                                                                   self.box.max_depth,
                                                                                   self.box.min_lat,
                                                                                   self.box.max_lat))

        cdo.vertmean(input=temp2, output=temp)
        nco.ncap2(input=temp, output=temp, options='-O -s "coslat[lat]=cos(lat[lat]*3.141592657/180.0)"')
        nco.ncwa(input=temp, output=temp, options='-w coslat -a lat')
        nco.ncks(input=temp, output=temp, options='-O -v vsftmyz,time')
        self.data_manager.send_file(temp, 'ocean', 'vsftmyz', self.startdate, self.member, self.chunk, box=self.box)