# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import numpy as np
import datetime

# import Apache OCW dependences
import ocw.data_source.local as local
from ocw.dataset import Bounds as Bounds
import ocw.dataset_processor as dsp
import ocw.metrics as metrics
import ocw.plotter as plotter
import ocw.utils as utils

# nClimDiv observation file
file_obs = 'nClimDiv/nClimDiv_tave_1895-2005.nc'

# CMIP5 simulations
model_file_path = 'CMIP5_historical'
dataset_name = ['BNU-ESM','CanESM2','CNRM-CM5','CSIRO-Mk3',
                'GISS-E2-H','HadCM3','inmcm4','MIROC-ESM',
                'MPI-ESM-LR','NorESM1-M']
nmodel = len(dataset_name) # number of CMIP5 simulations

# temporal boundary
start_date = datetime.datetime(1979, 12, 1)
end_date = datetime.datetime(2005, 8, 31)

nyear = 26                             

month_start = 6 # June
month_end = 8   # August

regions = []
regions.append(['WA', 'OR', 'ID'])
regions.append(['CA', 'NV', 'AZ', 'NM', 'UT','CO'])
regions.append(['MT', 'WY', 'ND', 'SD', 'NE'])
regions.append(['KS', 'TX', 'OK'])
regions.append(['MN', 'IA', 'MO', 'WI', 'IL', 'IN', 'MI','OH'])
regions.append(['ME', 'VT', 'NH', 'MA', 'NY', 'CT', 'RI','NJ','PA','WV','DE', 'MD'])
regions.append(['KY', 'VA', 'AR','AL', 'LA','MS', 'FL', 'GA','NC','SC', 'DC','TN'])

plotter.fill_US_states_with_color(regions, 'NCA_seven_regions', colors=True,
                                region_names=['NW','SW','NGP','SGP','MW','NE','SE'])

n_region = 7 # number of regions

# CONUS regional boundaries
NW_bounds = Bounds(boundary_type='us_states', 
                   us_states=regions[0])
SW_bounds = Bounds(boundary_type='us_states', 
                   us_states=regions[1])
NGP_bounds = Bounds(boundary_type='us_states', 
                   us_states=regions[2])
SGP_bounds = Bounds(boundary_type='us_states', 
                   us_states=regions[3])
MW_bounds = Bounds(boundary_type='us_states', 
                   us_states=regions[4])
NE_bounds = Bounds(boundary_type='us_states', 
                   us_states=regions[5])
SE_bounds = Bounds(boundary_type='us_states', 
                   us_states=regions[6])

regional_bounds = [NW_bounds, SW_bounds, NGP_bounds,
                   SGP_bounds, MW_bounds, NE_bounds, SE_bounds]

""" Load nClimDiv file into OCW Dataset """
obs_dataset = local.load_file(file_obs, variable_name='tave') 

""" Load CMIP5 simulations into a list of OCW Datasets"""
model_dataset = local.load_multiple_files(file_path=model_file_path, variable_name='tas',
                                          dataset_name=dataset_name, variable_unit='K') 

""" Temporal subset of obs_dataset """
obs_dataset_subset = dsp.temporal_slice(obs_dataset, 
                  start_time=start_date, end_time=end_date)
obs_dataset_season = dsp.temporal_subset(obs_dataset_subset, month_start, month_end,
                      average_each_year=True)

""" Temporal subset of model_dataset """
model_dataset_subset = [dsp.temporal_slice(dataset,start_time=start_date, end_time=end_date) 
                        for dataset in model_dataset]
model_dataset_season = [dsp.temporal_subset(dataset, month_start, month_end,
                      average_each_year=True) for dataset in model_dataset_subset]


""" Spatial subset of obs_dataset and generate time series """
obs_timeseries = np.zeros([nyear, n_region])   # region index 0-6: NW, SW, NGP, SGP, MW, NE, SE
model_timeseries = np.zeros([nmodel, nyear, n_region])

for iregion in np.arange(n_region):
    obs_timeseries[:, iregion] = utils.calc_time_series(
                         dsp.subset(obs_dataset_season, regional_bounds[iregion])) 
    for imodel in np.arange(nmodel):
        model_timeseries[imodel, :, iregion] = utils.calc_time_series(
                         dsp.subset(model_dataset_season[imodel], regional_bounds[iregion]))

year = np.arange(nyear)

regional_trends_obs = np.zeros(n_region)
regional_trends_obs_error = np.zeros(n_region)
regional_trends_model = np.zeros([nmodel, n_region])
regional_trends_model_error = np.zeros([nmodel, n_region])
regional_trends_ens = np.zeros(n_region)
regional_trends_ens_error = np.zeros(n_region)

for iregion in np.arange(n_region):
    regional_trends_obs[iregion], regional_trends_obs_error[iregion] = utils.calculate_temporal_trend_of_time_series(
                                                                       year, obs_timeseries[:,iregion])
    for imodel in np.arange(nmodel):
        regional_trends_model[imodel, iregion], regional_trends_model_error[iregion] = utils.calculate_temporal_trend_of_time_series(
                                                                                       year, model_timeseries[imodel, :, iregion])
    regional_trends_ens[iregion], regional_trends_ens_error[iregion] = utils.calculate_ensemble_temporal_trends(
                                                                       model_timeseries[:, :, iregion]) 

""" Generate plots """

plotter.fill_US_states_with_color(regions, 'nClimDiv_tave_trends_JJA_1980-2005', 
                                  values=regional_trends_obs,
                                  region_names=['%.3f' %(10*i) for i in regional_trends_obs])

plotter.fill_US_states_with_color(regions, 'CMIP5_ENS_tave_trends_JJA_1980-2005', 
                                  values=regional_trends_ens,
                                  region_names=['%.3f' %(10*i) for i in regional_trends_ens])

bias_ens = regional_trends_ens - regional_trends_obs
plotter.fill_US_states_with_color(regions, 'CMIP5_ENS_tave_trends_bias_from_nClimDiv_JJA_1980-2005', 
                                  values=bias_ens,
                                  region_names=['%.3f' %(10*i) for i in bias_ens])

obs_data = np.vstack([regional_trends_obs, regional_trends_obs_error])
ens_data = np.vstack([regional_trends_ens, regional_trends_ens_error])

plotter.draw_plot_to_compare_trends(obs_data, ens_data, regional_trends_model,
                                    fname='Trends_comparison_btn_CMIP5_and_nClimDiv',
                                    data_labels=['NW','SW','NGP','SGP','MW','NE','SE'],
                                    xlabel='NCA regions', ylabel='tas trend [K/year]')