#!/usr/bin/env python

# Copyright 2018 Earth Sciences Department, BSC-CNS
#
# This file is part of HERMESv3_GR.
#
# HERMESv3_GR is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# HERMESv3_GR is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with HERMESv3_GR. If not, see <http://www.gnu.org/licenses/>.


import os
import timeit
from warnings import warn as warning
import hermesv3_gr.config.settings as settings


class Masking(object):
    """
    Masking object to apply simple mask or factor mask.

    :param world_info: Path to the file that contains the ISO Codes and other relevant information.
    :type world_info: str

    :param factors_mask_values: List of the factor mask values.
    :type factors_mask_values: list

    :param regrid_mask_values: List of the mask values.
    :type regrid_mask_values: list

    :param grid: Grid.
    :type grid: Grid

    :param world_mask_file:
    :type world_mask_file: str
    """

    def __init__(self, world_info, factors_mask_values, regrid_mask_values, grid, world_mask_file=None,
                 countries_shapefile=None):
        from timezonefinder import TimezoneFinder

        st_time = timeit.default_timer()
        settings.write_log('\t\tCreating mask.', level=2)

        self.countries_shapefile = countries_shapefile
        self.adding = None
        self.world_info = world_info
        self.country_codes = self.get_country_codes()
        self.world_mask_file = world_mask_file
        self.factors_mask_values = self.parse_factor_values(factors_mask_values)
        self.regrid_mask_values = self.parse_masking_values(regrid_mask_values)
        self.regrid_mask = None
        self.scale_mask = None
        self.timezonefinder = TimezoneFinder()

        self.grid = grid

        settings.write_time('Masking', 'Init', timeit.default_timer() - st_time, level=3)

    def get_country_codes(self):
        """
        Get the country code information.

        :return: Dictionary of country codes.
        :rtype: dict
        """
        import pandas as pd

        st_time = timeit.default_timer()

        dataframe = pd.read_csv(self.world_info, sep=';')
        del dataframe['time_zone'], dataframe['time_zone_code']
        dataframe = dataframe.drop_duplicates().dropna()
        dataframe = dataframe.set_index('country_code_alpha')
        countries_dict = dataframe.to_dict()
        countries_dict = countries_dict['country_code']

        settings.write_time('Masking', 'get_country_codes', timeit.default_timer() - st_time, level=3)
        return countries_dict

    def create_country_iso(self, comm, netcdf_path):
        import numpy as np
        from hermesv3_gr.tools.netcdf_tools import netcdf2shp, extract_vars, gather_shapefile
        from hermesv3_gr.modules.writing.writer import Writer
        import geopandas as gpd

        st_time = timeit.default_timer()
        settings.write_log('\t\t\tCreating {0} file.'.format(self.world_mask_file), level=2)
        # output_path = os.path.join(output_dir, 'iso.nc')
        shapefile = netcdf2shp(comm, netcdf_path)

        if comm.Get_rank() == 0:
            countries = gpd.read_file(self.countries_shapefile)
            countries = countries.to_crs(shapefile.crs)
        else:
            countries = None
        countries = comm.bcast(countries, root=0)

        shapefile = gpd.sjoin(shapefile, countries[['M49', 'geometry']], how='left', op='intersects')
        shapefile = shapefile[~shapefile.index.duplicated(keep='first')]
        shapefile.fillna(0, inplace=True)
        shapefile = gather_shapefile(comm, shapefile)

        if settings.rank == 0:
            lat_o, lon_o = extract_vars(netcdf_path, ['lat', 'lon'])
            lon = np.array([lon_o['data']] * len(lat_o['data']))
            lat = np.array([lat_o['data']] * len(lon_o['data'])).T

            values = shapefile['M49'].astype(int).values
            data = [{
                'name': 'timezone_id',
                'units': '',
                'data': values.reshape((1, lat.shape[0], lat.shape[1])),
            }]

            Writer.write_netcdf(self.world_mask_file, lat, lon, data, regular_latlon=True)
        settings.comm.Barrier()

        settings.write_time('Masking', 'create_country_iso', timeit.default_timer() - st_time, level=3)

        return True

    def find_timezone(self, latitude, longitude):

        st_time = timeit.default_timer()

        if longitude < -180:
            longitude += 360
        elif longitude > +180:
            longitude -= 360

        tz = self.timezonefinder.timezone_at(lng=longitude, lat=latitude)

        settings.write_time('Masking', 'find_timezone', timeit.default_timer() - st_time, level=3)

        return tz

    def get_iso_code_from_tz(self, tz):
        import pandas as pd

        st_time = timeit.default_timer()

        zero_values = [None, ]
        if tz in zero_values:
            return 0

        df = pd.read_csv(self.world_info, sep=';')
        code = df.country_code[df.time_zone == tz].values

        settings.write_time('Masking', 'get_iso_code_from_tz', timeit.default_timer() - st_time, level=3)

        return code[0]

    def parse_factor_values(self, values):
        """

        :param values:
        :return:
        :rtype: dict
        """
        import re

        st_time = timeit.default_timer()

        if type(values) != str:
            return None
        values = list(map(str, re.split(' , |, | ,|,', values)))
        scale_dict = {}
        for element in values:
            element = list(map(str, re.split("{0}{0}|{0}".format(' '), element)))
            scale_dict[int(self.country_codes[element[0]])] = element[1]

        settings.write_log('\t\t\tApplying scaling factors for {0}.'.format(values), level=3)
        settings.write_time('Masking', 'parse_factor_values', timeit.default_timer() - st_time, level=3)

        return scale_dict

    def parse_masking_values(self, values):
        """

        :param values:
        :return:
        :rtype: list
        """
        import re

        st_time = timeit.default_timer()

        if type(values) != str:
            return None
        values = list(map(str, re.split(' , |, | ,|,| ', values)))
        if values[0] == '+':
            self.adding = True
        elif values[0] == '-':
            self.adding = False
        else:
            if len(values) > 0:
                settings.write_log('WARNING: Check the .err file to get more info. Ignoring mask')
                if settings.rank == 0:
                    warning("WARNING: The list of masking does not start with '+' or '-'. Ignoring mask.")
            return None
        code_list = []
        for country in values[1:]:
            code_list.append(int(self.country_codes[country]))

        if self.adding:
            settings.write_log("\t\t\tCreating mask to do {0} countries.".format(values[1:]), level=3)
        else:
            settings.write_log("\t\t\tCreating mask to avoid {0} countries.".format(values[1:]), level=3)
        settings.write_time('Masking', 'parse_masking_values', timeit.default_timer() - st_time, level=3)

        return code_list

    def check_regrid_mask(self, comm, input_file):

        if self.regrid_mask_values is not None:
            if not os.path.exists(self.world_mask_file):
                self.create_country_iso(comm, input_file)
            self.regrid_mask = self.custom_regrid_mask()
        if self.factors_mask_values is not None:
            if not os.path.exists(self.world_mask_file):
                self.create_country_iso(comm, input_file)
            self.scale_mask = self.custom_scale_mask()

    def custom_regrid_mask(self):
        import numpy as np
        from netCDF4 import Dataset

        st_time = timeit.default_timer()

        netcdf = Dataset(self.world_mask_file, mode='r')
        values = netcdf.variables['timezone_id'][:]
        netcdf.close()

        if self.adding:
            mask = np.zeros(values.shape)
            for code in self.regrid_mask_values:
                mask[values == code] = 1
        else:
            mask = np.ones(values.shape)
            for code in self.regrid_mask_values:
                mask[values == code] = 0

        settings.write_time('Masking', 'custom_regrid_mask', timeit.default_timer() - st_time, level=3)

        return mask

    def custom_scale_mask(self):
        import numpy as np
        from hermesv3_gr.tools.netcdf_tools import extract_vars

        st_time = timeit.default_timer()

        [values] = extract_vars(self.world_mask_file, ['timezone_id'])

        values = values['data']
        mask = np.ones(values.shape)
        for code, factor in self.factors_mask_values.items():
            mask[values == code] = factor

        settings.write_time('Masking', 'custom_scale_mask', timeit.default_timer() - st_time, level=3)

        return mask


if __name__ == '__main__':
    pass
