#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Created on Tue Jun 23 21:32:54 2020
# Author: Xiaodong Ming

"""
Rainfall
========

To do:
    To read, compute, and show rainfall data

-----------------

"""
import time
import numpy as np
import matplotlib.pyplot as plt	
import matplotlib.dates as mdates
from scipy import stats
from datetime import datetime
from datetime import timedelta
from hydro_raster.Raster import Raster
from . import indep_functions as indep_f
class Rainfall:
    """ a class to set rainfall data for hipims

    Attributes:

        time_s: 1-col array, time in seconds
        mask_header: dictionary showing mask georeference
        mask_dict: dict with two keys:'value' and 'index', providing int array
                showing rain source number and their index respectively
        rain_rate: numpy array m/s
        attrs: summary of the object
        subs_in: tuple of row and col number, provide subs of the mask array values
            inside the model domain. Only available when dem_ras is given
        start_date: start date and time of time zero in time_s
        time_dt: datetime format of time_s

    """
    def __init__(self, rain_mask, rain_source, source_sep=',', dem_ras=None):
        """initialize rainfall object with source file and mask file

        Args:
            rain_mask: str [filename of a Raster endswith .gz/asc/tif] numpy int array with the same shape with DEM array
                 a Raster object. If rain_mask is a scalar or array, dem_ras must be provided.           
            rain_source: numpy array the 1st column is time in seconds, 2nd to the end columns are rainfall rates in m/s. 
                str [filename of a csv file for rainfall source data]
            dem_ras: a Raster object for DEM
            source_sep: delimeter of the rain source file
            
        """
        self.set_mask(rain_mask, dem_ras)
        self.set_source(rain_source, source_sep)
        del rain_mask
        
    def set_mask(self, rain_mask, dem_ras=None):
        """Set rainfall mask from a scalar or a grid (object/file)

        if rain_mask is a scalar or array, dem_ras must be provided
        """
        if type(rain_mask) is str:
            obj_rain_mask = Raster(rain_mask)
            print(rain_mask+ ' read')
        elif hasattr(rain_mask, 'header'):
            obj_rain_mask = rain_mask
        else: # numpy array or a scalar, dem_ras must be provided
            rain_mask = np.array(rain_mask).astype('int32')
            dem_shape = dem_ras.array.shape
            if rain_mask.size > 1:
                if rain_mask.shape != dem_shape:
                    raise ValueError('The shape of rainfall_mask array '
                             'is not consistent with DEM')
            mask_array = np.zeros(dem_shape, dtype=np.int32)+rain_mask
            obj_rain_mask = Raster(array=mask_array, header=dem_ras.header)
        if hasattr(dem_ras, 'header'):
            # rain mask resample to the same shape with DEM
            self.mask_dict = indep_f._mask2dict(obj_rain_mask, dem_ras.header)
            self.mask_header = dem_ras.header
            self.subs_in = np.where(~np.isnan(dem_ras.array))
        else:
            self.mask_dict = indep_f._mask2dict(obj_rain_mask)
            self.mask_header = obj_rain_mask.header
    
    def set_source(self, rain_source, delimiter=','):
        """Set rainfall mask from a numpy array or a csv file
        """
        if type(rain_source) is str:
            print('reading '+rain_source)
            rain_source = np.loadtxt(rain_source, delimiter=delimiter)
        elif type(rain_source) is np.ndarray:
            pass
        else:
            raise IOError('rain_source must be either a filename or a numpy'
                          ' array')
#        _ = rp._check_rainfall_rate_values(rain_source, times_in_1st_col=True)
        self.time_s = rain_source[:, 0 ]
        self.rain_rate = rain_source[:, 1:]        
        self.get_attrs()

    def get_time_series(self, method='mean', rain_rate_valid=None):
        """ Plot time series of average rainfall rate inside the model domain   

        Args:
            method: 'mean'|'max','min','mean'method to calculate gridded rainfall 
                over the model domain
        """
        start = time.perf_counter()
        if rain_rate_valid is None:
            rain_rate_valid, _ = self.get_valid_rain_rate() # time-source id
        if method == 'mean':
            value_y = np.mean(rain_rate_valid, axis=1) # m/s
        elif method == 'max':
            value_y = np.max(rain_rate_valid, axis=1) # m/s
        elif method == 'min':
            value_y = np.min(rain_rate_valid, axis=1) # m/s
        elif method == 'median':
            value_y = np.median(rain_rate_valid, axis=1) # m/s
        elif method == 'sum':
            value_y = np.sum(rain_rate_valid, axis=1) # m/s
        else:
            raise ValueError('Cannot recognise the calculation method')
        value_y =  value_y*3600*1000 # mm/h
        time_series = np.c_[self.time_s, value_y]
        end = time.perf_counter()
        print('get_time_series took'+str(end-start))
        return time_series

    def get_spatial_map(self, method='sum'):
        """Get spatial rainfall map over time series

        Args:
            rain_mask_obj: asc file name or Raster object for rain mask
            cellsize: resample the rain_mask to a new grid (with larger cellsize)
            method: sum|mean caculate method for each cell, sum by time or mean by time 

        """
        # caculate rain source
        time_s = self.time_s
        rain_total = np.trapz(self.rain_rate, time_s, axis=0)*1000 #mm
        unit_str = 'mm/h'
        if method == 'sum':
            cell_rain = rain_total #mm
            unit_str = 'mm'
        elif method == 'mean':
            cell_rain = rain_total/(time_s.max()-time_s.min())*3600 #mm/h
        elif method == 'max':
            cell_rain = np.max(self.rain_rate, axis=0)*1000*3600 #mm/h
        elif method== 'min':
            cell_rain = np.min(self.rain_rate, axis=0)*1000*3600 #mm/h
        elif method== 'median':
            cell_rain = np.median(self.rain_rate, axis=0)*1000*3600 #mm/h
        else:
            raise ValueError('Cannot recognise the calculation method')    
        mask_values = self.mask_dict['value']
        rain_values = cell_rain[mask_values.astype('int32')]
        rain_dict = {'value':rain_values, 'index':self.mask_dict['index']}
        array_shape = (self.mask_header['nrows'], self.mask_header['ncols'])
        rain_array = indep_f._dict2grid(rain_dict, array_shape)
        return rain_array, unit_str
    
    def get_valid_rain_rate(self, unique=True):
        """Get a rain rate array over valid cells

        row is time axis, col is source id axis

        """
        array_shape = (self.mask_header['nrows'], self.mask_header['ncols'])
        mask_array = indep_f._dict2grid(self.mask_dict, array_shape)
        
        if hasattr(self, 'subs_in'):
            mask_valid = mask_array[self.subs_in]
        else:
            mask_valid = mask_array.flatten()
        mask_valid = mask_valid.astype('int32')
        if unique:
            rain_rate_valid = self.rain_rate[:, np.unique(mask_valid)]
        else:
            rain_rate_valid = self.rain_rate[:, mask_valid]
        return rain_rate_valid, mask_valid
    
    def set_start_date(self, datetime_str, fmt):
        """Set start date for the time series
        """
        self.start_date = datetime.strptime(datetime_str, fmt)
        time_delta = np.array([timedelta(seconds=i) for i in self.time_s])
        self.time_dt = self.start_date+time_delta
    
    def get_attrs(self):
        """ define rain_attr
        """
        rain_rate_valid, mask_valid = self.get_valid_rain_rate()
        time_s = self.time_s
        cellsize = self.mask_header['cellsize']
        temporal_res = (time_s.max()-time_s.min())/(time_s.size-1) # seconds
        source_uniq, source_counts = np.unique(mask_valid, return_counts=True)
        num_source = source_uniq.size
        mode_value = stats.mode(source_counts)[0]
        mode_value = np.array(mode_value).flatten()
        spatial_res = np.sqrt(mode_value[0])*cellsize
        cell_mean = np.mean(rain_rate_valid, axis=1)
        # total value along time axis for a cell
        rain_total = np.trapz(cell_mean, x=time_s, axis=0)*1000 # mm
        rain_mean = rain_total/(time_s.max()-time_s.min())*3600 # mm/h
        cell_max = np.max(rain_rate_valid, axis=1)*3600*1000
        attrs = {}
        attrs['num_source'] = num_source
        attrs['max'] = [np.max(cell_max).round(2), 'mm/h']
        attrs['sum'] = [rain_total.round(2), 'mm']
        attrs['average'] = [rain_mean.round(2), 'mm/h']
        attrs['spatial_res'] = [spatial_res.round(), 'm']
        attrs['temporal_res'] = [temporal_res.round(), 's']
        self.attrs = attrs
        return attrs
    
    def get_source_array(self):
        """Return a source array, first column is time_s, then continue with
        rain_rate in m/s
        """
        rain_source = np.c_[self.time_s, self.rain_rate]
        return rain_source
    
    def get_mask_array(self):
        """Return a mask array
        """
        array_shape = (self.mask_header['nrows'], self.mask_header['ncols'])
        mask_array = indep_f._dict2grid(self.mask_dict, array_shape)
        return mask_array
    
#%% Visualization	
    def plot_time_series(self, method='mean', dt_interval=24, 	
                         dt_format='%m-%d', title_str=None,	
                         **kwargs):	
        """ Plot time series of average rainfall rate inside the model domain   	
        method: 'mean'|'max','min','mean'method to calculate gridded rainfall 	
        over the model domain	
        """	
        value_y = self.get_time_series(method=method)
        value_y = value_y[:, -1]
        fig, ax = plt.subplots()        	
        if hasattr(self, 'time_dt'):	
            time_x = self.time_dt	
            ax.xaxis.set_major_locator(	
                    mdates.HourLocator(interval=dt_interval))	
            ax.xaxis.set_major_formatter(mdates.DateFormatter(dt_format))	
        else:	
            time_x = self.time_s	
        ax.plot(time_x, value_y, **kwargs)	
        ax.set_ylabel(method+' rainfall rate (mm/h)')	
        ax.grid(True)	
        if title_str is None:	
            title_str = method+' precipitation in the model domain'	
        ax.set_title(title_str)	
        plt.show()	
        return fig, ax

    def plot_rainfall_map(self, method='sum',cmap='YlGnBu', **kw):	
        """plot rainfall map within model domain	
        method: the way to calculate time series rainfall rate	
        """	
        rain_array, unit_str = self.get_spatial_map(method)	
        grid_obj = Raster(array=rain_array, header=self.mask_header)	
        fig, ax = grid_obj.mapshow(cax_str=unit_str, cmap=cmap, **kw)	
        return fig, ax	
