import warnings
warnings.simplefilter(action='ignore')
import sys
import itertools
from pathlib import Path
from datetime import datetime

import geopandas as gpd
import pandas as pd
import numpy as np

spanon = Path.cwd().parents[2].joinpath("Spatial-Anonymization/src/")
sys.path.append(str(spanon))
from spatial_anonymization import utils
from spatial_anonymization import point_displacement

# =============================
# BASE WORKING DIR
# ============================
DIR_DATA = Path.cwd().parents[1].joinpath('data', 'TZA')
DIR_SPATIAL_DATA = DIR_DATA.joinpath('spatial')
DIR_NPS_W4 = DIR_DATA.joinpath('surveys', 'NPS_w4')
DIR_OUTPUTS = DIR_NPS_W4.joinpath('geovars', 'puf')

# =================
# GEOVARS DATA
# =================
DIR_RFE = Path.cwd().parents[1].joinpath('data', 'RFE', 'dekad')

# =============================
# SURVEY COORDINATES 
# ============================
# Public/displaced GPS oordinates files
FILE_HH_PUB_GEOVARS = DIR_NPS_W4.joinpath('geovars', 'npsy4.ea.offset.dta')


# =============================
# ADMIN BOUNDARIES
# =============================
# use geoboundaries API for country bounds
FILE_ADM0 = DIR_SPATIAL_DATA.joinpath('geoBoundaries-TZA-ADM0-all',
                                     'geoBoundaries-TZA-ADM0.shp')

# ====================
# GENERAL
# =====================
DEBUG = False
# in case we need to download OSM data
GEOFABRICK_REGION  = 'africa'
COUNTRY_NAME = 'Tanzania'
COUNTRY_ISO = 'TZA'
FILENAME_SUFFIX = 'npsw4'
# When survey was conducted
SURVEY_CALENDAR_YR = 2015
DATA_EXTRACT_YEARS = ['2013-2014', '2014-2015']
# Year starts on July 1 (dekad-19) instead of
# Jan-1 (dekad01)
DEKAD_YEAR_START = 19
DEKADS = ['d{}'.format(str(i).zfill(2)) for i in range(DEKAD_YEAR_START, 37)] + \
['d{}'.format(str(i).zfill(2)) for i in range(1, DEKAD_YEAR_START)]
if DEBUG:
    for i in range(28):
        q = [DEKADS[i] for i in range(i, i+9)]
        assert len(q) == 9
        print('Quarter-{}: {}'.format(i+1, q))

# ==============
# BUFFER
# ==============
# whether to buffer points or
# extract values at point
BUFFER = False
BUFF_ZONE_DIST = 1000
    
# ====================
# INPUT DATA SPECS
# =====================
# Cols for input coordinates file
LAT = 'lat_modified' 
LON = 'lon_modified'
HH_ID = 'clusterid'

# =============================
# VARIABLE LABELS
# ============================
VAR_LABELS = {'anntot_avg': 'Average 12-month total rainfall for July-June',
    'wetQ_avg': 'Average total rainfall in wettest quarter within 12-month periods from July-June',
    'wetQ_avgstart': 'Average start of wettest quarter in dekads 1-36, where first dekad of July =1',
     'anntot_2013_2014': '12-month total rainfall in July-June, starting July 2013',
    'wetQ_2013_2014': 'Total rainfall in wettest quarter within 12- month periods starting July 2013',
     'wetQstart_2013_2014': 'Start of wettest quarter in dekads 1-36, where first dekad of July 2013=1',
     'anntot_2014_2015': '12-month total rainfall in July-June, starting July 2014',
     'wetQ_2014_2015': 'Total rainfall in wettest quarter within 12- month periods starting July 2014',
     'wetQstart_2014_2015': 'Start of wettest quarter in dekads 1-36, where first dekad of July 2014 =1'
    }

for name, label in VAR_LABELS.items():
    try:
        assert len(label) <= 80, 'LABEL TOO LONG'
    except AssertionError:
        print(name)

def get_year_and_tif_list(stop_year, year_start=1):
    """
    Retrieve filepath list and years
    """
    # ===================================
    # RETRIVE FILES BY CALENDAR YEAR
    # ===================================
    rfe_tifs = [i for i in DIR_RFE.iterdir() if i.suffix == ".tif"]
    yrs = {}
    for tif in rfe_tifs:
        yr = tif.parts[-1][4:8]
        if yrs.get(yr, None):
            items =  yrs[yr]
            items.append(tif)
            yrs[yr] = items
        else:
            yrs[yr] = [tif]
    
    # Keep only the years we need based on survey year
    # global variable se
    sorted_yrs = sorted(list(yrs.keys()))
    stop_idx = sorted_yrs.index(str(stop_year))
    calend_yrs = dict(sorted(yrs.items())[:stop_idx+1])
    
    # ===================================
    # DEAL WITH CROSS CALENDAR YEARS
    # ===================================
    if year_start != 1:
        cross_calend_yrs = {}
        files_lst = sorted(list(itertools.chain.from_iterable(list(calend_yrs.values()))))

        yr_start_id = int(list(calend_yrs.keys())[0])
        while year_start < len(files_lst):
            yr_end_id = yr_start_id + 1
            yr_id = '{}-{}'.format(yr_start_id, yr_end_id)
            yr_start_id += 1
            files = files_lst[year_start-1:year_start+35]
            cross_calend_yrs[yr_id] = files
            year_start += 36
        return cross_calend_yrs
    else:
        return calend_yrs

def get_wettest_quarter(df_yr, dekads=None):
    """
    Given one year data, get the quarter with maximum rains
    and start dekad of the quarter
    """
    dfq = pd.DataFrame(index=df_yr.index)
    quarter_start = {}
    if dekads:
        for i in range(28):
            q = [dekads[i] for i in range(i, i+9)]
            assert len(q) == 9
            colname = 'q{}'.format(i+1)
            quarter_start[colname] = int(q[0][1:])
            dfq[colname] =  df_yr[q].sum(axis=1)
    else:
        for i in range(1,29):
            q = ['d{}'.format(str(i).zfill(2)) for i in range(i, i+9)]
            assert len(q) == 9
            colname = 'q{}'.format(i)
            quarter_start[colname] = int(i)
            dfq[colname] =  df_yr[q].sum(axis=1)
    
    dfq['qrt_max_val'] = dfq.max(axis=1)
    dfq['qrt_max_id'] = dfq.idxmax(axis=1)
    dfq['qrt_max_start'] = dfq.qrt_max_id.apply(lambda x: quarter_start[x])
    
    return dfq[['qrt_max_val', 'qrt_max_id', 'qrt_max_start']]

def generate_yr_rfe_stats(tif_list, yr, gdf_coords, uniq_id):
    """
    Generates RFE stats for a single year as follows
    """
    df_lst = []
    for file in tif_list:
        colname = 'd{}'.format(file.parts[-1][8:10])
        df = utils.point_query_raster(in_raster=file, buffer=BUFFER,
                                        point_shp=gdf_coords, statcol_name=colname
                                        ,id_col=uniq_id,xcol=LON, ycol=LAT)
        df.set_index(uniq_id, inplace=True)
        df_lst.append(df)

    df = pd.concat(df_lst, axis=1)
    dekad_cols = sorted(df.columns)
    df = df[dekad_cols]

    # Annual total and max dekadal value
    df.loc[:,'total'] = df.sum(axis=1)
    df.loc[:,'max'] = df[dekad_cols].max(axis=1)

    # Wettest quarter
    dfq = get_wettest_quarter(df_yr=df, dekads=DEKADS)

    # merge 
    df2 = pd.concat([dfq, df[['total', 'max']]], axis=1)
    df2['yr'] = yr

    return df2

def generate_rfe_stats(yr_tif_dict, pts_shp, id_col):
    """
    Helper to loop through all years and run generate_yr_rfe_stats()
    """
    # ========================
    # CALCULATE YEARLY STATS
    # ========================
    df_yr_list = []
    # Loop years and calculate stats for that year
    cnt_yrs = 0
    for yr, files in yr_tif_dict.items():
        if len(files) != 36:
            print('SKIPPING {} BECAUSE WE DONT HAVE ALL DEKADS'.format(yr))
            continue
        df_yr = generate_yr_rfe_stats(files, yr, pts_shp, uniq_id=id_col)
        df_yr_list.append( df_yr)
        cnt_yrs += 1
        if DEBUG:
            if cnt_yrs == 5:
                break
    
    # concat into a single dataframe
    df = pd.concat(df_yr_list, axis=0)
    df = df[['yr',  'total', 'max', 'qrt_max_val', 'qrt_max_start']]
    assert len(df) == len(pts_shp)*cnt_yrs
    
    # ==================================
    # CALCULATE AVERAGES ACROSS YEARS
    # ==================================
    # Calculate the averages
    df_avg = pd.DataFrame(index=df.index.drop_duplicates())
    vars = ['total', 'max', 'qrt_max_val', 'qrt_max_start']
    print('Long average years: {}'.format(df.yr.unique()))
    for idx, row in df_avg.iterrows():
        data = df.loc[idx]
        for var in vars:
            new_colname = 'avg_{}'.format(var)
            avg = data[var].mean()
            df_avg.loc[idx, new_colname] = int(avg)
    
    # Add the averages to the original dataframe
    df2 = df.merge(df_avg, left_index=True, right_index=True)
    assert len(df2) == len(df)
    
    return df2

def generate_stats(target_yrs):
    """
    Helper function to put everything together and generate stats
    """
    # ===========================
    # LOAD COORDINATES
    # ===========================
    df_pub_coords = pd.read_stata(FILE_HH_PUB_GEOVARS, convert_categoricals=False)
    df_coords = df_pub_coords[[HH_ID, LAT, LON]]
    gdf_coords = gpd.GeoDataFrame(df_coords, 
                                      geometry=gpd.points_from_xy(df_coords[LON],df_coords[LAT]), crs=4326)
    
    # ===========================
    # GENERATE ALL YEARS STATS
    # ===========================
    start = datetime.now()
    files_yrs = get_year_and_tif_list(stop_year=SURVEY_CALENDAR_YR, year_start=DEKAD_YEAR_START)
    df = generate_rfe_stats(files_yrs, pts_shp=gdf_coords, id_col=HH_ID)
    end = datetime.now()
    time_taken = (end - start).total_seconds()/60
    print('Time taken: {}'.format(int(time_taken)))
    # ================================
    # EXTRACT DATA FOR TARGET YEARS
    # ================================
    df_target_yrs = df.query('yr in @target_yrs')
    required_vars = {'qrt_max_start': 'wetQstart', 'avg_total': 'anntot_avg', 
                     'avg_qrt_max_start': 'wetQ_avgstart', 'total': 'anntot',
                     'avg_qrt_max_val': 'wetQ_avg', 'qrt_max_val': 'wetQ'}
    df_target_yrs.rename(columns=required_vars, inplace=True)
    df_target_yrs = df_target_yrs[['yr', 'anntot', 'wetQ', 'wetQstart', 'anntot_avg','wetQ_avg', 'wetQ_avgstart']]

    # ===========================================
    # GENERATE YEAR BASED VARIABLES (E.G., 
    # ANNTOT_2011-2012,ANNTOT_2012-2013
    # ===========================================
    df_yr_based = df_target_yrs[['yr', 'anntot', 'wetQ', 'wetQstart']]

    df_tmp = pd.DataFrame(index=df_yr_based.index.drop_duplicates())
    yr_based_vars = df_yr_based.columns.difference(['yr'])
    
    # There is probably a shorter pivot approach but couldnt 
    # figure it out, so 
    for idx, row in df_tmp.iterrows():
        df_idx = df_yr_based.loc[idx]
        for yr in target_yrs:
            for var in yr_based_vars:
                val = df_idx.query('yr == @yr').iloc[0][var]
                new_col_name = '{}_{}'.format(var, yr.replace("-","_"))
                df_tmp.loc[idx, new_col_name] = val

    # ===========================================
    # PUT EVERYTHING BACK TOGETHER
    # ===========================================
    # We can drop everything else and remain with unique-locs
    df_target_yrs2 = df_target_yrs[~df_target_yrs.index.duplicated(keep='first')]
    df = df_target_yrs2.merge(df_tmp, left_index=True, right_index=True, how='inner')
    assert len(df) == len(gdf_coords)
    # Drop columns we dont need
    df.drop(columns=['yr', 'anntot', 'wetQ', 'wetQstart'], inplace=True)

    # ===========================================
    # SAVE THE OUTPUTS
    # =========================================== 
    out_vars = list(df.columns.difference([HH_ID, LON, LAT]))
    out_csv = DIR_OUTPUTS.joinpath('crop-season-{}.csv'.format(FILENAME_SUFFIX))
    df.to_csv(out_csv)
    
    
    # Save variable labels
    out_csv_labels = DIR_OUTPUTS.joinpath('crop-season-{}-labels.csv'.format(FILENAME_SUFFIX))
    df_var_labels = pd.DataFrame(list(VAR_LABELS.values()), index=list(VAR_LABELS.keys()))
    df_var_labels.index.name = 'var_name'
    df_var_labels.rename(columns={0:'label'}, inplace=True)
    df_var_labels.to_csv(out_csv_labels)
    
    print('-'*30)
    print('VARIABLES SAVED TO FILE BELOW')
    print('-'*30)
    print("/".join(out_csv.parts[6:]))
    
    return out_csv, out_vars

# Generate the variables
out_csv, out_vars = generate_stats(target_yrs=DATA_EXTRACT_YEARS)

# Only use describe for continuos variables 
print('-'*50)
print('       SUMMARY STATISTICS FOR CROP SEASON VARS')
print('-'*50)
stats = df[out_vars].describe().T
print(stats[['mean', 'min', '50%', '75%', 'max']])
print('-'*70)