#!/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 sys
import timeit
import hermesv3_gr.config.settings as settings
import numpy as np
from warnings import warn


class TemporalDistribution(object):
    """
    TemporalDistribution class that contains all the information for the temporal disaggregation.

    :param starting_date: Date of the first timestep.
    :type starting_date: datetime.datetime

    :param timestep_type: Relation between timesteps. It can be hourly, monthly or yearly.
    :type timestep_type: str

    :param timestep_num: Quantity of timesteps.
    :type timestep_num: int

    :param timestep_freq: Quantity of timestep_type between timesteps.
        eg: If timestep_type = hourly; timestep_freq = 2; The difference between time of each timestep is 2 hours.
    :type timestep_freq: int

    :param monthly_profile_path: Path to the file that contains all the monthly profiles.
    :type monthly_profile_path: str

    :param monthly_profile_id: ID of the monthly profile to use.
    :type monthly_profile_id: str

    :param daily_profile_path: Path to the file that contains all the daily profiles.
    :type weekly_profile_path: str

    :param daily_profile_id: ID of the daily profile to use.
    :type weekly_profile_id: str

    :param hourly_profile_path: Path to the file that contains all the hourly profiles.
    :type hourly_profile_path: str

    :param hourly_profile_id: ID of the hourly profile to use.
    :type hourly_profile_id: str

    :param world_info_path: Path to the file that contains the necessary information to do the NetCDF of timezones.
    :type world_info_path: str

    :param auxiliar_files_dir: Path to the directory where will be all the needed auxiliar files like the NetCDf of
        timezones.
    :type auxiliar_files_dir: str
    """
    def __init__(self, starting_date, timestep_type, timestep_num, timestep_freq, monthly_profile_path,
                 monthly_profile_id, weekly_profile_path, weekly_profile_id, daily_profile_path, daily_profile_id,
                 hourly_profile_path, hourly_profile_id, world_info_path, auxiliar_files_dir, grid, first_time=False):
        from timezonefinder import TimezoneFinder

        import pandas as pd

        st_time = timeit.default_timer()
        settings.write_log('\t\tInitializing Temporal.', level=2)

        self.grid = grid

        if not first_time:
            self.date_array = self.caclulate_date_array(starting_date, timestep_type, timestep_num, timestep_freq)

            if timestep_type in ['daily', 'hourly']:
                self.daily_profile = self.get_daily_profile(daily_profile_id, daily_profile_path)
            else:
                self.daily_profile = None

            if self.daily_profile is None:
                if timestep_type in ['monthly', 'daily', 'hourly']:
                    self.monthly_profile = self.get_monthly_profile(monthly_profile_id, monthly_profile_path)
                else:
                    self.monthly_profile = None

                if timestep_type in ['daily', 'hourly']:
                    self.weekly_profile = self.get_weekly_profile(weekly_profile_id, weekly_profile_path)
                else:
                    self.weekly_profile = None
            else:
                self.monthly_profile = None
                self.weekly_profile = None
                if monthly_profile_id is not None or weekly_profile_id is not None:
                    warn('WARNING: Daily profile is set. Monthly or Weekly profiles will be ignored.')

            if timestep_type in ['hourly']:
                self.hourly_profile = self.get_hourly_profile(hourly_profile_id, hourly_profile_path)
            else:
                self.hourly_profile = None

        self.world_info = world_info_path
        self.netcdf_timezones = os.path.join(auxiliar_files_dir, 'timezones.nc')

        self.world_info_df = pd.read_csv(self.world_info, sep=';')
        self.tf = TimezoneFinder()

        if not self.is_created_netcdf_timezones():
            settings.write_log("\t\tTimezones netCDF is not created. Lets try to create it.", level=1)
            self.create_netcdf_timezones(grid)
        self.timezones_array = self.calculate_timezones()

        settings.write_time('TemporalDistribution', 'Init', timeit.default_timer() - st_time, level=2)

    @staticmethod
    def caclulate_date_array(st_date, time_step_type, time_step_num, time_step_freq):
        from datetime import timedelta
        from calendar import monthrange, isleap

        if time_step_type == 'hourly':
            end_date = st_date + (time_step_num - 1) * timedelta(hours=time_step_freq)
        elif time_step_type == 'daily':
            end_date = st_date + (time_step_num - 1) * timedelta(hours=time_step_freq * 24)
        elif time_step_type == 'monthly':
            delta_year = (time_step_num - 1) * time_step_freq // 12
            delta_month = (time_step_num - 1) * time_step_freq % 12
            end_date = st_date.replace(year=st_date.year + delta_year,
                                       month=st_date.month + delta_month)
        elif time_step_type == 'yearly':
            delta_year = (time_step_num - 1) * time_step_freq
            end_date = st_date.replace(year=st_date.year + delta_year)
        else:
            end_date = st_date

        date_aux = st_date
        date_array = []
        while date_aux <= end_date:
            date_array.append(date_aux)  # 3600 seconds per hour

            if time_step_type == 'hourly':
                delta = timedelta(hours=time_step_freq)
            elif time_step_type == 'daily':
                delta = timedelta(hours=time_step_freq * 24)
            elif time_step_type == 'monthly':
                days = monthrange(date_aux.year, date_aux.month)[1]
                delta = timedelta(hours=days * 24)
            elif time_step_type == 'yearly':
                if isleap(date_aux.year):
                    delta = timedelta(hours=366 * 24)
                else:
                    delta = timedelta(hours=365 * 24)
            else:
                delta = None

            date_aux = date_aux + delta
        return date_array

    def get_month_array(self, date_array):
        from datetime import date

        month_array = [date(year=x.year, month=x.month, day=1) for x in date_array]
        month_array = list(set(month_array))

        return month_array

    def get_csv_temporal_profile(self, profile_id, profile_path):
        import pandas as pd
        profile = pd.read_csv(profile_path)
        profile = profile.loc[profile['ID'] == profile_id, :].drop(columns='ID')
        if len(profile) is 0:
            settings.write_log('ERROR: Check the .err file to get more info.')
            if settings.rank == 0:
                raise AttributeError('ERROR: Profile ID {0} is not in the {1} file.'.format(profile_id, profile_path))
            sys.exit(1)
        elif len(profile) > 1:
            settings.write_log('ERROR: Check the .err file to get more info.')
            if settings.rank == 0:
                raise AttributeError('ERROR: Profile ID {0} is repeated {2} times in the {1} file.'.format(
                    profile_id, profile_path, len(profile)))
            sys.exit(1)
        return profile.reset_index(drop=True)

    def get_gridded_temporal_profile(self, profile_value, profile_path):
        from netCDF4 import Dataset
        nc_in = Dataset(profile_path, mode='r')
        if self.grid.grid_type in ['rotated','rotated_nested']:
            lat_name = 'rlat'
            lon_name = 'rlon'
        elif self.grid.grid_type in ['mercator', 'lcc']:
            lat_name = 'x'
            lon_name = 'y'
        else:
            lat_name = 'lat'
            lon_name = 'lon'
        try:
            if nc_in.dimensions[lat_name].size * nc_in.dimensions[lon_name].size != self.grid.full_shape[-2] * \
                    self.grid.full_shape[-1]:
                settings.write_log('ERROR: Check the .err file to get more info.')
                if settings.rank == 0:
                    raise IOError('ERROR: The temporal gridded profile {0} '.format(profile_path) +
                                  'is not in the destiny resolution like {0}.'.format(self.grid.coords_netcdf_file))
                sys.exit(1)

        except KeyError:
            if nc_in.dimensions['latitude'].size * nc_in.dimensions['longitude'].size != self.grid.full_shape[-2] * \
                    self.grid.full_shape[-1]:
                settings.write_log('ERROR: Check the .err file to get more info.')
                if settings.rank == 0:
                    raise IOError('ERROR: The temporal gridded profile {0} '.format(profile_path) +
                                  'is not in the destiny resolution like {0}.'.format(self.grid.coords_netcdf_file))

        try:
            var = nc_in.variables[profile_value][:, self.grid.x_lower_bound:self.grid.x_upper_bound,
                                                 self.grid.y_lower_bound:self.grid.y_upper_bound]
        except KeyError:
            settings.write_log('ERROR: Check the .err file to get more info.')
            if settings.rank == 0:
                raise AttributeError('ERROR: The variable {0} is not in the {1} file.'.format(
                    profile_value, profile_path))
            sys.exit(1)
        var = var.reshape((var.shape[0], var.shape[1] * var.shape[2]))
        return var

    def get_daily_profile(self, profile_id, profile_path):
        """
        Function to read the daily profile.
        Return options:
            - None -> If no daily profile set.
            - dict -> 365 (or 365) dict keys with the Julian day as key.
            - numpy.ndarray -> 3D array with the gridded profiles.

        :param profile_id: That parameter can be None if no daily profile to set; A path to the gridded profile;
            Or a string (length=4) with the ID of the profile (it have to appear in the profile path).

        :param profile_path: Path to the CSV file that contains all the daily profiles. The sum of the profile have to
            be the amount of days of the year to simulate.
        :type profile_path: str

        :return: Daily profile (numpy.ndarray if gridded or dict for the rest)
        :rtype: None, numpy.ndarray, dict
        """
        import calendar

        if calendar.isleap(self.date_array[0].year):
            days_of_year = 366
        else:
            days_of_year = 365

        try:
            if np.isnan(profile_id):
                profile_id = None
        except TypeError:
            profile_id = profile_id

        # No Daily profile
        if profile_id is None:
            profile = None
        # Gridded monthly profile
        elif os.path.exists(profile_id):
            profile = self.get_gridded_temporal_profile('Fday', profile_id)
            profile_sum = profile.sum() / profile.shape[-1]
            if not calendar.isleap(self.date_array[0].year):
                profile = np.append(profile, profile[-1, :].reshape((1, profile.shape[-1])), axis=0)
        # Plane daily profile
        elif len(profile_id) is 4:
            profile = self.get_csv_temporal_profile(profile_id, profile_path)
            profile_sum = profile.T[0].sum()
            profile.columns = profile.columns.astype(int)
            profile = profile.to_dict('records')[0]
            if not calendar.isleap(self.date_array[0].year):
                profile[366] = profile[365]
        else:
            settings.write_log('ERROR: Check the .err file to get more info.')
            if settings.rank == 0:
                raise AttributeError('ERROR: An error has occurred in the profile {0}.'.format(profile_id))
            sys.exit(1)

        # CHECK if it is normalized
        if profile is not None and round(days_of_year, 2) != round(profile_sum, 2):
            settings.write_log('ERROR: Check the .err file to get more info.')
            if settings.rank == 0:
                raise AttributeError('ERROR: The sum of the temporal weight factors is not {0} (is {1}).'.format(
                    days_of_year, profile_sum))
            sys.exit(1)

        return profile

    def get_monthly_profile(self, profile_id, profile_path):
        """
        Function to read the monthly profile.
        Return options:
            - None -> If no monthly profile set.
            - dict -> 12 dict keys with the month number (from 1 to 12) as key.
            - numpy.ndarray -> 3D array with the gridded profiles.

        :param profile_id: That parameter can be None if no monthly profile to set; A path to the gridded profile;
            Or a string (length=4) with the ID of the profile (it have to appear in the profile path).

        :param profile_path: Path to the CSV file that contains all the monthly profiles. The sum of the profile have to
            be 12.
        :type profile_path: str

        :return: Monthly profile (numpy.ndarray if gridded or dict for the rest)
        :rtype: None, numpy.ndarray, dict
        """
        try:
            if np.isnan(profile_id):
                profile_id = None
        except TypeError:
            profile_id = profile_id

        # No Daily profile
        if profile_id is None:
            profile = None
        # Gridded monthly profile
        elif os.path.exists(profile_id):
            profile = self.get_gridded_temporal_profile('Fmonth', profile_id)
            profile_sum = profile.sum() / profile.shape[-1]
        # Plane daily profile
        elif len(profile_id) is 4:
            profile = self.get_csv_temporal_profile(profile_id, profile_path)
            profile_sum = profile.T[0].sum()
            profile.rename(columns={
                'January': 1, 'February': 2, 'March': 3, 'April': 4, 'May': 5, 'June': 6, 'July': 7, 'August': 8,
                'September': 9, 'October': 10, 'November': 11, 'December': 12}, inplace=True)
            profile = profile.to_dict('records')[0]
        else:
            settings.write_log('ERROR: Check the .err file to get more info.')
            if settings.rank == 0:
                raise AttributeError('ERROR: An error has occurred in the profile {0}.'.format(profile_id))
            sys.exit(1)

        # CHECK if it is normalized
        if profile is not None and round(12, 2) != round(profile_sum, 2):
            settings.write_log('ERROR: Check the .err file to get more info.')
            if settings.rank == 0:
                raise AttributeError('ERROR: The sum of the temporal weight factors is not {0} (is {1}).'.format(
                    12, profile_sum))
            sys.exit(1)

        return profile

    def rebalance_weekly_profile(self, date, profile):
        weekdays_count = self.calculate_weekdays(date)
        if isinstance(profile, dict):
            factor = self.calculate_weekday_factor(profile, weekdays_count)
            for dict_key in profile.keys():
                profile[dict_key] = profile[dict_key] + factor
        elif isinstance(profile, np.ndarray):
            # Gridded
            factor = self.calculate_weekday_gridded_factor(profile, weekdays_count)
            for weekday in range(7):
                profile[weekday, :] = profile[weekday, :] + factor
        else:
            settings.write_log('ERROR: Check the .err file to get more info.')
            if settings.rank == 0:
                raise TypeError("ERROR: Profile type '{0}' not implemented".format(type(profile)))
            sys.exit(1)
        return profile

    def get_weekly_profile(self, profile_id, profile_path):
        try:
            if np.isnan(profile_id):
                return None
        except TypeError:
            profile_id = profile_id

        profile = {0: None}
        for month_date in self.get_month_array(self.date_array):
            # No Daily profile
            if profile_id is None:
                return None
            # Gridded monthly profile
            elif os.path.exists(profile_id):
                profile_aux = self.get_gridded_temporal_profile('Fweek', profile_id)
                profile_sum = profile_aux.sum() / profile_aux.shape[-1]
            # Plane daily profile
            elif len(profile_id) is 4:
                profile_aux = self.get_csv_temporal_profile(profile_id, profile_path)
                profile_sum = profile_aux.T[0].sum()
                profile_aux.rename(columns={'Monday': 0, 'Tuesday': 1, 'Wednesday': 2, 'Thursday': 3, 'Friday': 4,
                                            'Saturday': 5, 'Sunday': 6}, inplace=True)
                profile_aux = profile_aux.to_dict('records')[0]
            else:
                settings.write_log('ERROR: Check the .err file to get more info.')
                if settings.rank == 0:
                    raise AttributeError('ERROR: An error has occurred in the profile {0}.'.format(profile_id))
                sys.exit(1)

            # CHECK if it is normalized
            if profile_aux is not None and round(7, 2) != round(profile_sum, 2):
                settings.write_log('ERROR: Check the .err file to get more info.')
                if settings.rank == 0:
                    raise AttributeError('ERROR: The sum of the temporal weight factors is not {0} (is {1}).'.format(
                        7, profile_sum))
                sys.exit(1)
            else:
                if profile[0] is None:
                    profile[0] = profile_aux.copy()
                profile[month_date.month] = self.rebalance_weekly_profile(month_date, profile_aux)
        return profile

    @staticmethod
    def parse_hourly_profile_id(profile_id):
        import re

        dict_aux = {}
        try:
            list_aux = list(map(str, re.split(' , | ,|, |,| ', profile_id)))
        except TypeError:
            return None
        if len(list_aux) is 1:
            dict_aux['weekday'] = profile_id
            dict_aux['saturday'] = profile_id
            dict_aux['sunday'] = profile_id
        else:
            for element in list_aux:
                key_value_list = list(map(str, re.split(':| :|: | : |=| =|= | = ', element)))
                dict_aux[key_value_list[0]] = key_value_list[1]

        return dict_aux

    def get_hourly_profile(self, profile_id, profile_path):
        """
        Function to read the hourly profile.
        Return options:
            - None -> If no hourly profile set.
            - dict -> 24 dict keys with the month number (from 0 to 23) as key.
            - numpy.ndarray -> 3D array with the gridded profiles.

        :param profile_id: That parameter can be None if no hourly profile to set; A path to the gridded profile;
            Or a string (length=4) with the ID of the profile (it have to appear in the profile path).

        :param profile_path: Path to the CSV file that contains all the hourly profiles. The sum of the profile have to
            be 24.
        :type profile_path: str

        :return: Hourly profile (numpy.ndarray if gridded or dict for the rest)
        :rtype: None, numpy.ndarray, dict
        """
        try:
            if np.isnan(profile_id):
                profile = None
        except TypeError:
            profile = self.parse_hourly_profile_id(profile_id)

        if profile is not None:
            for profile_type, profile_id_aux in profile.items():
                # Gridded monthly profile
                if os.path.exists(profile_id_aux):
                    profile_aux = self.get_gridded_temporal_profile('Fhour', profile_id_aux)
                    profile_sum = profile_aux.sum() / profile_aux.shape[-1]
                # Plane daily profile
                elif len(profile_id_aux) is 4:
                    profile_aux = self.get_csv_temporal_profile(profile_id_aux, profile_path)
                    profile_sum = profile_aux.T[0].sum()
                    profile_aux.columns = profile_aux.columns.astype(int)
                    profile_aux = profile_aux.to_dict('records')[0]
                else:
                    settings.write_log('ERROR: Check the .err file to get more info.')
                    if settings.rank == 0:
                        raise AttributeError('ERROR: An error has occurred in the profile {0}.'.format(profile_id_aux))
                    sys.exit(1)

                # CHECK if it is normalized
                if profile_aux is not None and round(24, 2) != round(profile_sum, 2):
                    settings.write_log('ERROR: Check the .err file to get more info.')
                    if settings.rank == 0:
                        raise AttributeError('ERROR: The sum of the temporal weight factors ' +
                                             '{0} is not {1} (is {2}).'.format(profile_id_aux, 24, profile_sum))
                    sys.exit(1)
                else:
                    profile[profile_type] = profile_aux

        return profile

    def get_tz_from_id(self, tz_id):
        """
        Extract the timezone (string format) for the given id (int).

        :param tz_id: ID of the timezone.
        :type tz_id: int

        :return: Timezone
        :rtype: str
        """
        tz = self.world_info_df.time_zone[self.world_info_df.time_zone_code == tz_id].values

        return tz[0]

    def get_id_from_tz(self, tz):
        """
        Extract the id (int) for the given timezone (string format).

        :param tz: Timezone of the ID.
        :type tz: str

        :return: ID
        :rtype: int
        """
        tz_id = self.world_info_df.time_zone_code[self.world_info_df.time_zone == tz].values

        try:
            tz_id = tz_id[0]
        except IndexError:
            settings.write_log('ERROR: Check the .err file to get more info.')
            if settings.rank == 0:
                raise ValueError("ERROR: The timezone '{0}' is not in the {1} file. ".format(tz, self.world_info) +
                                 "Please add it as a new line with an unique time_zone_code " +
                                 "and the corresponding country and country_code.")
            sys.exit(1)

        return tz_id

    @staticmethod
    def parse_tz(timezone):
        """
        Parse the timezone (string format).

        It is needed because some libraries have more timezones than others and it
        tries to simplify setting the strange ones into the nearest common one.
        Examples:
            'America/Punta_Arenas': 'America/Santiago',
            'Europe/Astrakhan': 'Europe/Moscow',
            'Asia/Atyrau': 'Asia/Aqtau',
            'Asia/Barnaul': 'Asia/Almaty',
            'Europe/Saratov': 'Europe/Moscow',
            'Europe/Ulyanovsk': 'Europe/Moscow',
            'Europe/Kirov': 'Europe/Moscow',
            'Asia/Tomsk': 'Asia/Novokuznetsk',
            'America/Fort_Nelson': 'America/Vancouver'

        :param timezone: Not parsed timezone.
        :type timezone: str

        :return: Parsed timezone
        :rtype: str
        """
        tz_dict = {
            'America/Punta_Arenas': 'America/Santiago',
            'Europe/Astrakhan': 'Europe/Moscow',
            'Asia/Atyrau': 'Asia/Aqtau',
            'Asia/Barnaul': 'Asia/Almaty',
            'Europe/Saratov': 'Europe/Moscow',
            'Europe/Ulyanovsk': 'Europe/Moscow',
            'Europe/Kirov': 'Europe/Moscow',
            'Asia/Tomsk': 'Asia/Novokuznetsk',
            'America/Fort_Nelson': 'America/Vancouver',
            'Asia/Famagusta': 'Asia/Nicosia',
        }

        if timezone in iter(tz_dict.keys()):
            timezone = tz_dict[timezone]

        return timezone

    def find_closest_timezone(self, latitude, longitude):
        """
        Find the closest timezone for the given coordinates.

        :param latitude: Latitude coordinate to find timezone.
        :type latitude: float

        :param longitude: Longitude coordinate fo find the timezone.
        :type longitude: float

        :return: Nearest timezone of the given coordinates.
        :rtype: str
        """
        st_time = timeit.default_timer()

        degrees = 0
        timezone = None
        while timezone is None:
            timezone = self.tf.closest_timezone_at(lng=longitude, lat=latitude, delta_degree=degrees)
            degrees += 1

        settings.write_time('TemporalDistribution', 'find_closest_timezone', timeit.default_timer() - st_time, level=3)

        return timezone

    def is_created_netcdf_timezones(self):
        """
        Check if the NetCDF of timezones is created

        :return: True if it is already created.
        :rtype: bool
        """
        return os.path.exists(self.netcdf_timezones)

    def create_netcdf_timezones(self, grid):
        """
        Create a NetCDF with the timezones in the resolution of the given grid.

        :param grid: Grid object with the coordinates.
        :type grid: Grid

        :return: True if it is created.
        :rtype: bool
        """
        from hermesv3_gr.tools.netcdf_tools import write_netcdf

        st_time = timeit.default_timer()
        settings.write_log("\t\tCreating {0} file.".format(self.netcdf_timezones), level=2)

        lat, lon = grid.get_coordinates_2d()
        total_lat = settings.comm.gather(lat, root=0)
        total_lon = settings.comm.gather(lon, root=0)

        dst_var = []

        num = 0
        points = list(zip(lat.flatten(), lon.flatten()))

        for lat_aux, lon_aux in points:
            num += 1
            settings.write_log("\t\t\tlat:{0}, lon:{1} ({2}/{3})".format(lat_aux, lon_aux, num, len(points)), level=3)
            timezone = self.find_closest_timezone(lat_aux, lon_aux)
            tz_id = self.get_id_from_tz(timezone)
            dst_var.append(tz_id)
        dst_var = np.array(dst_var)
        dst_var = dst_var.reshape((1,) + lat.shape)
        dst_var = settings.comm.gather(dst_var, root=0)
        if settings.rank == 0:
            total_lat = np.concatenate(total_lat, axis=1)
            total_lon = np.concatenate(total_lon, axis=1)
            dst_var = np.concatenate(dst_var, axis=2)
            data = [{'name': 'timezone_id', 'units': '', 'data': dst_var}]

            write_netcdf(self.netcdf_timezones, total_lat, total_lon, data, regular_latlon=True)
        settings.comm.Barrier()

        settings.write_time('TemporalDistribution', 'create_netcdf_timezones', timeit.default_timer() - st_time,
                            level=2)

        return True

    # def read_gridded_profile(self, path, value):
    #     # TODO Documentation
    #     """
    #
    #     :param path:
    #     :param value:
    #     :return:
    #     """
    #     from netCDF4 import Dataset
    #
    #     st_time = timeit.default_timer()
    #
    #     settings.write_log('\t\t\tGetting gridded temporal monthly profile from {0} .'.format(path), level=3)
    #
    #     nc_in = Dataset(path)
    #     profile = nc_in.variables[value][:, self.grid.x_lower_bound:self.grid.x_upper_bound,
    #                                      self.grid.y_lower_bound:self.grid.y_upper_bound]
    #     nc_in.close()
    #
    #     profile[profile <= 0] = 1
    #
    #     settings.write_time('TemporalDistribution', 'read_gridded_profile', timeit.default_timer() - st_time, level=3)
    #
    #     return profile

    def calculate_timezones(self):
        """
        Calculate the timezones ID's from the NetCDF and convert them to the timezone (str).

        :return: Array with the timezone of each cell.
        :rtype: numpy.chararray
        """
        from netCDF4 import Dataset

        st_time = timeit.default_timer()

        nc_in = Dataset(self.netcdf_timezones)
        timezones = nc_in.variables['timezone_id'][:, self.grid.x_lower_bound:self.grid.x_upper_bound,
                                                   self.grid.y_lower_bound:self.grid.y_upper_bound].astype(int)

        nc_in.close()
        tz_list = np.chararray(timezones.shape, itemsize=32)
        for id_aux in range(timezones.min(), timezones.max() + 1):
            try:
                timezone = self.get_tz_from_id(id_aux)
                tz_list[timezones == id_aux] = timezone
            except IndexError:
                pass

        settings.write_time('TemporalDistribution', 'calculate_timezones', timeit.default_timer() - st_time, level=3)

        return tz_list

    def calculate_2d_temporal_factors(self, date_aux):
        """
        Calculate the temporal factor to correct the input data of the given date for each cell.

        :param date_aux: Date of the current timestep.
        :type date_aux: datetime.datetime

        :return: 2D array with the factors to correct the input data to the date of this timestep.
        :rtype: numpy.array
        """
        from datetime import date
        import pytz
        import pandas as pd

        st_time = timeit.default_timer()

        df = pd.DataFrame(self.timezones_array.flatten(), columns=['tz'])

        df['local'] = pd.to_datetime(date_aux, utc=True)
        try:
            df['local'] = df.groupby('tz')['local'].apply(
                lambda x: x.dt.tz_convert(x.name.decode("utf-8")).dt.tz_localize(None))
        except pytz.exceptions.UnknownTimeZoneError:
            df['local'] = df.groupby('tz')['utc'].apply(
                lambda x: x.dt.tz_convert(self.parse_tz(x.name.decode("utf-8"))).dt.tz_localize(None))

        # ===== HOURLY PROFILES =====
        df['weekday'] = df['local'].dt.weekday
        df['hour'] = df['local'].dt.hour
        if self.hourly_profile is not None:
            if isinstance(self.hourly_profile, dict):
                # WEEKDAY
                weekday_profile = self.hourly_profile['weekday']
                if isinstance(weekday_profile, dict):
                    df['weekday_factor'] = df['hour'].map(weekday_profile)
                else:
                    for hour in range(24):
                        df.loc[df['hour'] == hour, 'weekday_factor'] = weekday_profile[
                            hour, df[df['hour'] == hour].index]
                # SATURDAY
                saturday_profile = self.hourly_profile['saturday']
                if isinstance(saturday_profile, dict):
                    df['saturday_factor'] = df['hour'].map(saturday_profile)
                else:
                    for hour in range(24):
                        df.loc[df['hour'] == hour, 'saturday_factor'] = saturday_profile[
                            hour, df[df['hour'] == hour].index]
                # SUNDAY
                sunday_profile = self.hourly_profile['sunday']
                if isinstance(sunday_profile, dict):
                    df['sunday_factor'] = df['hour'].map(sunday_profile)
                else:
                    for hour in range(24):
                        df.loc[df['hour'] == hour, 'sunday_factor'] = sunday_profile[
                            hour, df[df['hour'] == hour].index]

                # Selecting profile type depending on the weekday
                df.loc[df['weekday'] <= 4, 'hourly_factor'] = df['weekday_factor'][df['weekday'] <= 4].values
                df.loc[df['weekday'] == 5, 'hourly_factor'] = df['saturday_factor'][df['weekday'] == 5].values
                df.loc[df['weekday'] == 6, 'hourly_factor'] = df['sunday_factor'][df['weekday'] == 6].values

                df.drop(columns=['weekday_factor', 'saturday_factor', 'sunday_factor'], inplace=True)
        else:
            df['hourly_factor'] = 1

        df.drop(columns=['hour'], inplace=True)

        # ===== WEEKLY PROFILES =====
        df['month'] = df['local'].dt.month
        if self.weekly_profile is not None:
            for month in np.unique(df['month']):
                if month not in self.weekly_profile:
                    self.weekly_profile[month] = self.rebalance_weekly_profile(
                        date(year=df.loc[df['month'] == month, 'local'].dt.year.values[0], month=month, day=1),
                        self.weekly_profile[0])

                if isinstance(self.weekly_profile[month], dict):
                    df.loc[df['month'] == month, 'weekly_factor'] = df['weekday'].map(self.weekly_profile[month])
                else:
                    for weekday in range(7):
                        df.loc[df['weekday'] == weekday, 'weekly_factor'] = self.weekly_profile[month][
                            weekday, df[df['weekday'] == weekday].index]

        else:
            df['weekly_factor'] = 1

        df.drop(columns=['weekday'], inplace=True)

        # ===== MONTHLY PROFILES =====
        if self.monthly_profile is not None:
            if isinstance(self.monthly_profile, dict):
                df['monthly_factor'] = df['month'].map(self.monthly_profile)
            else:
                for month in np.unique(df['month']):
                    df.loc[df['month'] == month, 'monthly_factor'] = self.monthly_profile[
                        month - 1, df[df['month'] == month].index]
        else:
            df['monthly_factor'] = 1

        df.drop(columns=['month'], inplace=True)

        # ===== DAILY PROFILES =====
        df['day'] = df['local'].dt.dayofyear
        if self.daily_profile is not None:
            if isinstance(self.daily_profile, dict):
                df['daily_factor'] = df['day'].map(self.daily_profile)
            else:
                for day in np.unique(df['day']):
                    df.loc[df['day'] == day, 'daily_factor'] = self.daily_profile[
                        day - 1, df[df['day'] == day].index]
        else:
            df['daily_factor'] = 1

        df.drop(columns=['day'], inplace=True)

        df['factor'] = df['monthly_factor'] * df['weekly_factor'] * df['daily_factor'] * df['hourly_factor']

        df.drop(columns=['monthly_factor', 'weekly_factor', 'daily_factor', 'hourly_factor'], inplace=True)

        factors = np.array(df['factor'].values).reshape((self.timezones_array.shape[1], self.timezones_array.shape[2]))
        del df

        settings.write_time('TemporalDistribution', 'calculate_2d_temporal_factors', timeit.default_timer() - st_time,
                            level=3)

        return factors

    def calculate_3d_temporal_factors(self):
        """
        Calculate the temporal factor to correct the input data of the given date for each cell.

        :return: 3D array with the factors to correct the input data to the date of this timestep.
        :rtype: numpy.array
        """
        st_time = timeit.default_timer()
        settings.write_log("\tCalculating temporal factors.", level=2)

        factors = []
        # date = self.starting_date

        for i, date_aux in enumerate(self.date_array):
            settings.write_log("\t\t{0} temporal factor ({1}/{2}).".format(
                date_aux.strftime('%Y/%m/%d %H:%M:%S'), i+1, len(self.date_array)), level=3)

            factors.append(self.calculate_2d_temporal_factors(date_aux))

            # self.hours_since.append((date_aux - self.date_array[0]).seconds / 3600 + date_aux.days * 24)

        factors = np.array(factors)

        settings.write_time('TemporalDistribution', 'calculate_3d_temporal_factors', timeit.default_timer() - st_time,
                            level=3)
        return factors

    @staticmethod
    def calculate_weekday_factor(profile, weekdays):
        # TODO Documentation
        """
        Operate with all the days of the month to get the sum of daily factors of the full month.

        :param profile:
        :param weekdays:
        :return:
        """
        st_time = timeit.default_timer()

        weekdays_factors = 0
        num_days = 0
        for weekday in range(7):
            weekdays_factors += profile[weekday] * weekdays[weekday]
            num_days += weekdays[weekday]

        settings.write_time('TemporalDistribution', 'calculate_weekday_factor',
                            timeit.default_timer() - st_time, level=3)

        return (num_days - weekdays_factors) / num_days

    @staticmethod
    def calculate_weekday_gridded_factor(profile, weekdays):
        # TODO Documentation
        """
        Operate with all the days of the month to get the sum of daily factors of the full month.

        :param profile:
        :param weekdays:
        :return:
        """
        st_time = timeit.default_timer()

        weekdays_factors = np.zeros((profile.shape[-1]))
        num_days = 0
        for weekday in range(7):
            weekdays_factors += profile[weekday, :] * weekdays[weekday]
            num_days += weekdays[weekday]

        factor = (num_days - weekdays_factors) / num_days

        settings.write_time('TemporalDistribution', 'calculate_weekday_gridded_factor',
                            timeit.default_timer() - st_time, level=3)

        return factor

    @staticmethod
    def calculate_weekdays(date):
        # TODO Documentation
        """

        :param date:
        :return:
        """
        from calendar import monthrange, weekday, MONDAY, TUESDAY, WEDNESDAY, THURSDAY, FRIDAY, SATURDAY, SUNDAY

        st_time = timeit.default_timer()

        weekdays = [MONDAY, TUESDAY, WEDNESDAY, THURSDAY, FRIDAY, SATURDAY, SUNDAY]
        days = [weekday(date.year, date.month, d + 1) for d in range(monthrange(date.year, date.month)[1])]
        weekdays_dict = {}
        count = 0
        for day in weekdays:
            weekdays_dict[count] = days.count(day)

            count += 1

        settings.write_time('TemporalDistribution', 'calculate_weekdays', timeit.default_timer() - st_time, level=3)
        return weekdays_dict

    # @staticmethod
    # def get_temporal_monthly_profile(profile_path, profile_id):
    #     """
    #     Extract the monthly profile of the given ID in a dictionary format.
    #
    #     The month (1 to 12) is the key (int) and the value (float) is the factor.
    #
    #     :param profile_path: Path to the file that contains all the monthly profiles.
    #     :type profile_path: str
    #
    #     :param profile_id: ID of the monthly profile to use.
    #     :type profile_id: str
    #
    #     :return: Monthly profile where the month (1 to 12) is the key (int) and the value (float) is the factor.
    #     :rtype: dict
    #     """
    #     import pandas as pd
    #
    #     st_time = timeit.default_timer()
    #
    #     settings.write_log("\t\t\tGetting temporal monthly profile id '{0}' from {1} .".format(
    #         profile_id, profile_path), level=3)
    #
    #     if profile_id is not None:
    #         df = pd.read_csv(profile_path)
    #         try:
    #             profile = df.loc[df[df.TP_M == profile_id].index[0]].to_dict()
    #         except IndexError:
    #             settings.write_log('ERROR: Check the .err file to get more info.')
    #             if settings.rank == 0:
    #                 raise AttributeError('ERROR: Monthly profile ID {0} is not in the {1} file.'.format(
    #                     profile_id, profile_path))
    #             sys.exit(1)
    #         profile.pop('TP_M', None)
    #         profile = {int(k): float(v) for k, v in profile.items()}
    #     else:
    #         profile = None
    #
    #     settings.write_time('TemporalDistribution', 'get_temporal_monthly_profile', timeit.default_timer() - st_time,
    #                         level=2)
    #
    #     return profile

    @staticmethod
    def calculate_delta_hours(st_date, time_step_type, time_step_num, time_step_freq):
        # TODO Documentation
        """

        :param st_date:
        :param time_step_type:
        :param time_step_num:
        :param time_step_freq:
        :return:
        """
        from datetime import timedelta
        from calendar import monthrange, isleap

        st_time = timeit.default_timer()

        settings.write_log('Calculating time array of {0} time steps starting from {1}.'.format(
            time_step_num, st_date.strftime('%Y/%m/%d %H:%M:%S')))

        if time_step_type == 'hourly':
            end_date = st_date + (time_step_num - 1) * timedelta(hours=time_step_freq)
        elif time_step_type == 'daily':
            end_date = st_date + (time_step_num - 1) * timedelta(hours=time_step_freq * 24)
        elif time_step_type == 'monthly':
            delta_year = (time_step_num - 1) * time_step_freq // 12
            delta_month = (time_step_num - 1) * time_step_freq % 12
            end_date = st_date.replace(year=st_date.year + delta_year,
                                       month=st_date.month + delta_month)
        elif time_step_type == 'yearly':
            delta_year = (time_step_num - 1) * time_step_freq
            end_date = st_date.replace(year=st_date.year + delta_year)
        else:
            end_date = st_date

        date_aux = st_date
        hours_since = []
        while date_aux <= end_date:
            d = date_aux - st_date
            hours_since.append(d.seconds / 3600 + d.days * 24)  # 3600 seconds per hour

            if time_step_type == 'hourly':
                delta = timedelta(hours=time_step_freq)
            elif time_step_type == 'daily':
                delta = timedelta(hours=time_step_freq * 24)
            elif time_step_type == 'monthly':
                days = monthrange(date_aux.year, date_aux.month)[1]
                delta = timedelta(hours=days * 24)
            elif time_step_type == 'yearly':
                if isleap(date_aux.year):
                    delta = timedelta(hours=366 * 24)
                else:
                    delta = timedelta(hours=365 * 24)
            else:
                delta = None

            date_aux = date_aux + delta

        settings.write_time('TemporalDistribution', 'calculate_delta_hours', timeit.default_timer() - st_time, level=2)

        return hours_since


if __name__ == '__main__':
    pass
