import numpy as np
from autosubmit.config.log import Log

from earthdiagnostics.models import Models
from earthdiagnostics.diagnostic import Diagnostic
from earthdiagnostics.utils import Utils, TempFile


class Gyres(Diagnostic):
    """
    Compute the intensity of the subtropical and subpolar gyres

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

    :created: October 2013
    :last modified: June 2016

    """
    def __init__(self, data_manager, startdate, member, chunk, nemo_version):
        Diagnostic.__init__(self, data_manager)
        self.startdate = startdate
        self.member = member
        self.chunk = chunk
        self.nemo_version = nemo_version
        self.required_vars = ['vsftbarot']
        self.generated_vars = ['gyres']

    def __str__(self):
        return 'Gyres Startdate: {0} Member: {1} Chunk: {2} '.format(self.startdate, self.member,
                                                                     self.chunk)

    @classmethod
    def generate_jobs(cls, diags, options):
        job_list = list()
        for startdate in diags.startdates:
            for member in diags.members:
                for chunk in range(1, diags.chunks + 1):
                    job_list.append(Gyres(diags.data_manager, startdate, member, chunk, diags.nemo_version))
        return job_list

    # noinspection PyPep8Naming
    def compute(self):
        if self.nemo_version in [Models.ECEARTH_2_3_O1L42, Models.ECEARTH_3_0_O1L46,
                                 Models.NEMO_3_2_O1L42, Models.NEMO_3_3_O1L46,
                                 Models.NEMOVAR_O1L42]:

            subpolNAtl = [230, 275, 215, 245]
            subpolNPac = [70, 145, 195, 235]
            subtropNPac = [45, 175, 165, 220]
            subtropNAtl = [195, 275, 175, 225]
            subtropSPac = [70, 205, 120, 145]
            subtropSAtl = [235, 300, 120, 145]
            subtropInd = [320, 30, 110, 180]
            ACC = [1, 361, 1, 65]

        elif self in [Models.ECEARTH_3_0_O25L46, Models.ECEARTH_3_0_O25L75,
                      Models.GLORYS2_V1_O25L75]:
            raise Exception("Option gyres not available yet for {0}".format(self.nemo_version))
        else:
            raise Exception("Input grid {0} not recognized".format(self.nemo_version))

        output = TempFile.get()
        vsftbarot_file = self.data_manager.get_file('ocean', 'vsftbarot', self.startdate, self.member, self.chunk)

        handler_original = Utils.openCdf(vsftbarot_file)
        self.var_vsftbarot = handler_original.variables['vsftbarot']
        handler = Utils.openCdf(output, 'w')
        handler.createDimension('time', handler_original.variables['time'].shape[0])
        handler.createDimension('region', 8)
        Utils.copy_variable(handler_original, handler, 'time')
        var_region = handler.createVariable('region', str, 'region')

        var_gyre = handler.createVariable('gyre', 'f', ('time', 'region'), fill_value=0.0)
        var_gyre.valid_max = 2e8
        var_gyre.valid_min = 0.0
        var_gyre.short_name = 'gyre'
        var_gyre.long_name = 'gyre'
        var_gyre.units = 'm^3/s'

        var_region[0] = 'subpolNAtl'
        var_gyre[:, 0] = self._gyre(subpolNAtl, True)
        Log.debug('subpolNAtl: {0}', var_gyre[:, 0])

        var_region[1] = 'subpolNPac'
        var_gyre[:, 1] = self._gyre(subpolNPac, True)
        Log.debug('subpolNPac: {0}', var_gyre[:, 1])

        var_region[2] = 'subtropNPac'
        var_gyre[:, 2] = self._gyre(subtropNPac)
        Log.debug('subtropNPac: {0}', var_gyre[:, 2])

        var_region[3] = 'subtropSPac'
        var_gyre[:, 3] = self._gyre(subtropSPac)
        Log.debug('subtropSPac: {0}', var_gyre[:, 3])

        var_region[4] = 'subtropNAtl'
        var_gyre[:, 4] = self._gyre(subtropNAtl)
        Log.debug('subtropNAtl: {0}', var_gyre[:, 4])

        var_region[5] = 'subtropSAtl'
        var_gyre[:, 5] = self._gyre(subtropSAtl)
        Log.debug('subtropSAtl: {0}', var_gyre[:, 5])

        var_region[6] = 'subtropInd'
        var_gyre[:, 6] = self._gyre(subtropInd)
        Log.debug('subtropInd: {0}', var_gyre[:, 6])

        var_region[7] = 'ACC'
        var_gyre[:, 7] = self._gyre(ACC)
        Log.debug('ACC: {0}', var_gyre[:, 7])

        handler.close()
        handler_original.close()
        self.data_manager.send_file(output, 'ocean', 'gyre', self.startdate, self.member, self.chunk)
        Log.info('Finished gyres for startdate {0}, member {1}, chunk {2}', self.startdate, self.member, self.chunk)

    def _gyre(self, site, invert=False):
        if invert:
            return np.min(self._extract_section(site), (1, 2)) * -1
        else:
            return np.max(self._extract_section(site), (1, 2))

    def _extract_section(self, site):
        if site[2] <= site[3]:
            if site[0] <= site[1]:
                return self.var_vsftbarot[:, site[2] - 1:site[3] - 1, site[0] - 1:site[1] - 1]
            else:
                return np.concatenate((self.var_vsftbarot[:, site[2] - 1:site[3] - 1, site[0]-1:],
                                      self.var_vsftbarot[:, site[2] - 1:site[3] - 1, :site[1] - 1]), axis=2)

        else:
            if site[0] <= site[1]:
                return np.concatenate((self.var_vsftbarot[:, site[2] - 1:, site[0] - 1: site[1] - 1],
                                      self.var_vsftbarot[:, :site[3] - 1, site[0] - 1: site[1] - 1]), axis=1)
            else:
                temp = np.concatenate((self.var_vsftbarot[:, site[2] - 1:, :],
                                      self.var_vsftbarot[:, :site[3] - 1, :]), axis=1)
                return np.concatenate((temp[:, :, site[0] - 1:],
                                      temp[:, :, :site[1] - 1]), axis=2)