# coding=utf-8
import numpy as np
from autosubmit.config.log import Log

from earthdiagnostics import cdftools
from earthdiagnostics.constants import Basins
from earthdiagnostics.diagnostic import Diagnostic
from earthdiagnostics.utils import Utils, TempFile


class Moc(Diagnostic):
    """
    Compute the MOC for oceanic basins

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

    :created: March 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
    """

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

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

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

        :param diags: Diagnostics manager class
        :type diags: Diags
        :param options: None
        :type options: list[str]
        :return:
        """
        if len(options) > 1:
            raise Exception('The MOC diagnostic has no options')
        job_list = list()
        for startdate, member, chunk in diags.exp_manager.get_chunk_list():
            job_list.append(Moc(diags.data_manager, startdate, member, chunk))
        return job_list

    def compute(self):
        """
        Runs the diagnostic
        """
        temp = TempFile.get()

        input_file = self.data_manager.get_file('ocean', 'vo', self.startdate, self.member, self.chunk)

        Log.debug('Computing MOC')
        cdftools.run('cdfmoc', input=input_file, output=temp)
        Utils.nco.ncks(input=input_file, output=temp, options='-A -v lev')
        Utils.convert2netcdf4(temp)

        Log.debug('Reformatting variables')
        handler = Utils.openCdf(temp)

        handler.createDimension('basin', 5)
        handler.createVariable('basin', str, 'basin')
        handler.variables['basin'][:] = np.array([Basins.Global.fullname, Basins.Atlantic.fullname,
                                                  Basins.Pacific.fullname, Basins.IndoPacific.fullname,
                                                  Basins.Indian.fullname], dtype=object)
        example = handler.variables['zomsfglo']
        # noinspection PyProtectedMember
        moc = handler.createVariable('vsftmyz', example.datatype,
                                     ('time', 'lev', 'i', 'j', 'basin'),
                                     fill_value=example._FillValue)

        moc.units = example.units
        moc.add_offset = example.add_offset
        moc.scale_factor = example.scale_factor

        moc[:, :, :, :, 0] = handler.variables['zomsfglo'][:]
        moc[:, :, :, :, 1] = handler.variables['zomsfatl'][:]
        moc[:, :, :, :, 2] = handler.variables['zomsfpac'][:]
        moc[:, :, :, :, 3] = handler.variables['zomsfinp'][:]
        moc[:, :, :, :, 4] = handler.variables['zomsfind'][:]

        handler.close()

        Utils.nco.ncks(input=temp, output=temp,
                       options='-O -x -v zomsfglo,zomsfatl,zomsfpac,zomsfinp,zomsfind,zomsfinp0')
        Utils.setminmax(temp, 'vsftmyz')

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