variable.py

# coding=utf-8
import csv
import json

import os
from autosubmit.config.log import Log

from earthdiagnostics.constants import Basins
from earthdiagnostics.domain import Domain


class VariableJsonException(Exception):
    pass


class SingletonType(type):
    def __call__(cls, *args):
        try:
            return cls.__instance
        except AttributeError:
            cls.__instance = super(SingletonType, cls).__call__(*args)
            return cls.__instance


class VariableManager(object):
    __metaclass__ = SingletonType

    def __init__(self):
        self._cmor_tables_folder = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'cmor_tables')
        self._aliases_folder = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'variable_alias')
        self._dict_variables = {}

    def get_variable(self, original_name, silent=False):
        """
        Returns the cmor variable instance given a variable name

        :param original_name: original variable's name
        :type original_name: str
        :param silent: if True, omits log warning when variable is not found
        :type silent: bool
        :return: CMOR variable
        :rtype: Variable
        """
        try:
            return self._dict_aliases[original_name.lower()][1]
        except KeyError:
            if not silent:
                Log.warning('Variable {0} is not defined in the CMOR table. Please add it'.format(original_name))
            return None

    def get_variable_and_alias(self, original_name, silent=False):
        """
        Returns the cmor variable instance given a variable name

        :param original_name: original variable's name
        :type original_name: str
        :param silent: if True, omits log warning when variable is not found
        :type silent: bool
        :return: CMOR variable
        :rtype: Variable
        """
        try:
            return self._dict_aliases[original_name.lower()]
        except KeyError:
            if not silent:
                Log.warning('Variable {0} is not defined in the CMOR table. Please add it'.format(original_name))
            return None

    def load_variables(self, table_name):
        """
        Loads the  CMOR csv and creates the variables dictionary
        """
        self._dict_variables = dict()
        self._load_variable_list(table_name)
        self._load_known_aliases()
        self._construct_aliases_dict()

    def _load_variable_list(self, table_name):
        csv_path = self._get_csv_path(table_name)
        if os.path.isfile(csv_path):
            self._load_file('default')
            self._load_file(table_name)
            return

        json_folder = self._get_json_folder(table_name)
        if os.path.isdir(json_folder):
            self._load_json(json_folder)

    def _get_csv_path(self, table_name):
        csv_table_path = os.path.join(self._cmor_tables_folder, '{0}.csv'.format(table_name))
        return csv_table_path

    def _get_json_folder(self, table_name):
        json_folder = os.path.join(self._cmor_tables_folder, '{0}/Tables'.format(table_name))
        return json_folder

    def _load_file(self, csv_table_path):
        with open(self._get_csv_path(csv_table_path), 'rb') as csvfile:
            reader = csv.reader(csvfile, dialect='excel')
            for line in reader:
                if line[0] == 'Variable':
                    continue

                var = Variable()
                var.parse_csv(line)
                if not var.short_name:
                    continue

                self._dict_variables[var.short_name.lower()] = var

    def _load_json(self, json_folder):
        for file_name in os.listdir(json_folder):
            if file_name in ('CMIP6_grids.json', 'CMIP6_formula_terms.json'):
                continue
            json_data = open(os.path.join(json_folder, file_name)).read()
            data = json.loads(json_data)
            if 'variable_entry' in data:
                Log.debug('Parsing file {0}'.format(file_name))
                self._load_json_variables(data['variable_entry'])

    def _load_json_variables(self, json_data):
        for short_name in json_data.keys():
            variable = Variable()
            try:
                variable.parse_json(json_data[short_name], short_name)
                self._dict_variables[variable.short_name.lower()] = variable
            except VariableJsonException:
                Log.error('Could not read variable {0}'.format(short_name))

    def _load_known_aliases(self):
        with open(self._get_aliases_csv_path(), 'rb') as csvfile:
            reader = csv.reader(csvfile, dialect='excel')
            for line in reader:
                if line[0] == 'Aliases':
                    continue

                aliases = line[0].split(':')
                if line[1].lower() in self._dict_variables:
                    cmor_var = self._dict_variables[line[1].lower()]
                else:
                    cmor_vars = []
                    for alias in aliases:
                        if alias.lower() in self._dict_variables:
                            cmor_vars.append(self._dict_variables[alias.lower()])
                    if len(cmor_vars) == 0:
                        Log.error('Aliases {0} could not be mapped to any variable'.format(aliases))
                        continue
                    elif len(cmor_vars) > 1:
                        Log.error('Aliases {0} can be be mapped to multiple variables '
                                  '{1}'.format(aliases, cmor_vars))
                        continue
                    cmor_var = cmor_vars[0]

                for alias in aliases:
                    if alias != cmor_var.short_name and alias in self._dict_variables:
                        Log.error('Alias {0} for variable {1} is already a different '
                                  'variable!'.format(alias, cmor_var.short_name))
                        continue
                    alias_object = VariableAlias(alias)
                    if line[2]:
                        alias_object.basin = Basins.parse(line[2])
                    if line[3]:
                        alias_object.grid = line[3]
                    cmor_var.known_aliases.append(alias_object)

    def _get_aliases_csv_path(self):
        csv_table_path = os.path.join(self._aliases_folder, 'default.csv')
        return csv_table_path

    def _construct_aliases_dict(self):
        self._dict_aliases = {}
        for cmor_var_name in self._dict_variables:
            cmor_var = self._dict_variables[cmor_var_name]
            if cmor_var_name not in cmor_var.known_aliases:
                cmor_var.known_aliases.append(VariableAlias(cmor_var_name))
            Log.debug('Variable: {0} Alias: {1}'.format(cmor_var_name,
                                                        ', '.join([str(alias) for alias in cmor_var.known_aliases])))
            for alias in cmor_var.known_aliases:
                self._dict_aliases[alias.alias] = (alias, cmor_var)


class Variable(object):
    """
    Class to characterize a CMOR variable. It also contains the static method to make the match between thje original
    name and the standard name. Requires data _convetion to be available in cmor_tables to work.
    """

    def __str__(self):
        return '{0} ({1})'.format(self.standard_name, self.short_name)

    def __init__(self):
        self.short_name = None
        self.standard_name = None
        self.long_name = None
        self.units = None
        self.valid_min = None
        self.valid_max = None
        self.grid = None
        self.known_aliases = []

    def parse_json(self, json_var, key):

        if 'out_name' in json_var:
            self.short_name = json_var['out_name']
        else:
            raise VariableJsonException('Variable has no out name defined'.format(key))
        self.standard_name = json_var['standard_name']
        self.long_name = json_var['long_name']

        domain = json_var['modeling_realm'].split(' ')
        if len(domain) > 1:
            Log.warning('Multiple modeling realms assigned to variable {0}: {1}. '
                        'We wil use first ({1[0]}) as domain'.format(self.short_name, domain))
        if not domain[0]:
            Log.warning('Variable {0} has no modeling realm defined'.format(self.short_name))
        else:
            self.domain = Domain(domain[0])

        self.valid_min = json_var['valid_min']
        self.valid_max = json_var['valid_max']
        self.units = json_var['units']

    def parse_csv(self, var_line):
        self.short_name = var_line[1].strip()
        self.standard_name = var_line[2].strip()
        self.long_name = var_line[3].strip()
        self.domain = Domain(var_line[4].strip())
        self.basin = Basins.parse(var_line[5])
        self.units = var_line[6].strip()
        self.valid_min = var_line[7].strip()
        self.valid_max = var_line[8].strip()
        self.grid = var_line[9].strip()


class VariableAlias(object):
    """
    Class to characterize a CMOR variable. It also contains the static method to make the match between thje original
    name and the standard name. Requires data _convetion to be available in cmor_tables to work.
    """

    def __init__(self, alias):
        self.alias = alias
        self.basin = None
        self.grid = None

    def __str__(self):
        string = self.alias
        if self.basin:
            string += ' Basin: {0}'.format(self.basin)
        if self.grid:
            string += ' Grid: {0}'.format(self.grid)
        return string




class VarType(object):
    MEAN = 1
    STATISTIC = 2

    @staticmethod
    def to_str(vartype):
        if vartype == VarType.MEAN:
            return 'mean'
        elif vartype == VarType.STATISTIC:
            return 'statistics'
        else:
            raise ValueError('Variable type {0} not supported'.format(vartype))