import pandas as pd
import datetime as dt
import seaborn as sns
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import matplotlib.ticker as ticker
import matplotlib as mpl
from matplotlib.colors import ListedColormap

REFERENCE_START = 1961
REFERENCE_END = 1990

FIGSIZE = (10, 8)
MAROON = '#550000'
BROWN = '#7A4300'
ORANGE = '#D07D18'
LIGHTBLUE = '#77A6CF'
BLUE = '#080043'
GRAY = '#86858A'

def prepare_data(city):
    '''
    Loads the data for the given city from "open-meteo-city.csv"
    Preprocesses the data
    '''
    
    # --- Read data ---
    df = pd.read_csv('open-meteo-{}.csv'.format(city), skiprows=3)

    # --- Prepare the dataset ---
    df.rename(
        columns={
            'temperature_2m (°C)' : 'temperature', 'time' : 'datetime'
            }, 
        inplace=True
        )
    df['datetime'] = pd.to_datetime(df['datetime'])
    df['date'] = df['datetime'].dt.date
    df['day'] = df['datetime'].dt.day
    df['month'] = df['datetime'].dt.month
    df['year'] = df['datetime'].dt.year

    # --- Delete data for February 29 ---
    df = df[(df['day'] != 29) | (df['month'] != 2)] 

    # --- Compute daily max and mean temperatures ---
    daily = df.groupby(
        ['year', 'month', 'day']
        )['temperature'].agg(
            ['max', 'mean']
            ).reset_index()

    # --- Prepare the reference values ---
    reference_daily = daily[
                        (daily['year'] >= REFERENCE_START)
                        & (daily['year'] <= REFERENCE_END)
                        ]
    reference = reference_daily.groupby(
        ['month','day']
        ).agg(
            {'max': ['mean', 'std'], 'mean': ['mean', 'std']}
            ).reset_index()
    # change the two-level headings to one-level
    reference.columns = [
        '_'.join(filter(None, col)).strip() 
        if isinstance(col, tuple) 
        else col 
        for col in reference.columns
        ] 

    # --- Return ---
    return df, daily, reference

MIN_EXC = 5 # minimal temperature excess
MIN_DAYS = 5 # minimal number of consecutive days

def heatwave_compare_data(city, how):
    ''' 
    Prepares a dataset used for heatwave plotting
    Uses heatwave definition with the difference: 
            MIN_EXC (= 5°C) if how = 'fixed'
            std if how = 'std'
    '''

    df, daily, ref = prepare_data(city)
    compare = daily.merge(
        ref, 
        how='left', 
        on=['month','day'], 
        suffixes=('_year', '_ref')
        )
    compare.rename(
        columns={
            'max': 'year_max', # daily max
            'mean' : 'year_mean', # daily mean
            'max_mean' : 'ref_max_mean', # mean over reference period of daily max
            'max_std' : 'ref_max_std', # std over reference period of daily max
            'mean_mean' : 'ref_mean_mean', # mean over reference period of daily mean
            'mean_std' : 'ref_mean_std' # std over reference period of daily max
            }, 
        inplace=True
        )
    if how == 'fixed':
        compare['hot'] = compare['ref_max_mean'
                ] + MIN_EXC < compare['year_max']
    elif how == 'std':
        compare['hot'] = compare['ref_max_mean'
                ] + compare['ref_max_std'] < compare['year_max'] 
    else:  
        raise ValueError('how must be "fixed" or "std".')
        
    compare['heatwave_end'] = (
          compare['hot'].rolling(MIN_DAYS).min() == 1
          ) 
    # However, we miss days 1-4 of each heatwave
    compare['heatwave'] = (
          compare['heatwave_end'][::-1].rolling(MIN_DAYS).max()[::-1] == 1
          )
    # Last four days was a heatwave in Milan 
    # (which we erased in the previous step)
    if city == 'Milan': 
        compare.loc[compare.index[-4:], 'heatwave'] = True
    
    # Only store the temperature if the day was part of a heatwave
    compare['heatwave_temp'] = None
    compare.loc[
          compare['heatwave'], 'heatwave_temp'
          ] = compare['year_max']
    
    return compare

def heatwave_days_in_year(df, year):
    ''' 
    Counts the number of days in the given year that satisfy 
    the definition of a heatwave.
    '''
    
    df_year = df[df['year'] == year]
    return df_year['heatwave'].sum()

def heatwave_count_plot(city, how):
    '''
    Creates a barplot with number of heatwave days per year 
    for the given city.
    
    Uses heatwave definition with the difference: 
            MIN_EXC (= 5°C) if how = 'fixed'
            std if how = 'std'
    '''
    
    compare = heatwave_compare_data(city, how)

    heatwave_days = []
    for year in range(1941, 2027):
        heatwave_days.append(
            heatwave_days_in_year(compare, year)
            )
    
    heatwave_days_reference = []
    for year in range(REFERENCE_START, REFERENCE_END+1):
        heatwave_days_reference.append(
            heatwave_days_in_year(compare, year)
            )


    heatwave_days_reference_mean = np.mean(
        heatwave_days_reference
        )

    fig = plt.figure(figsize=FIGSIZE)
    xaxis = range(1941, 2027)
    ax = sns.barplot(x=xaxis, y=heatwave_days, color=MAROON)
    ax.axhline(
        y=heatwave_days_reference_mean, 
        color=BLUE, linestyle='--', 
        label = 'Average of {start} – {end}'.format(
            start=REFERENCE_START, end=REFERENCE_END
            )
        )
    
    if how == 'fixed':
        method = str(MIN_EXC)+'°C'
    else:
        method = 'std'
    ax.set_title(
        'Number of days in year that count as a heatwave '
        'in {city} ({method})'.format(city=city, method=method),
        fontweight="bold", size=14, pad=10
        )
    ax.set_xticks(range(0, len(xaxis), 10))
    ax.yaxis.set_major_locator(ticker.MultipleLocator(10))
    ax.grid(True, axis='y', color='gray', alpha=0.3)
    plt.legend()
    plt.show()
    plt.clf()

def heatwave_days_plot(city, how, years):
    '''
    Plots the temperatures during heatwaves for the given 
    city and given years (list).
    
    Uses heatwave definition with the difference: 
            MIN_EXC (= 5°C) if how = 'fixed'
            std if how = 'std'
    '''

    # --- load and prepare data ---
    df, daily, ref = prepare_data(city)
    compare = heatwave_compare_data(city, how)

    # --- PLOT ---
    fig = plt.figure(figsize=FIGSIZE)
    days = range(1,366)

    # --- X axis: ticks at 1st of each month ---
    month_starts = []
    month_labels = [''] * 13 # hide tick labels
    month_names = [
        'Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 
        'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec',''
        ]

    base = dt.date(2026, 1, 1)
    for m in range(1, 13):
        d = dt.date(2026, m, 1)
        day_of_year = (d - base).days
        month_starts.append(day_of_year)

    month_starts.append(365) # Add Dec 31 as final boundary

    # --- plot ---
    ax = plt.subplot()
    plt.plot(
        ref['max_mean'], 
        label='Average of {start} – {end}'.format(
            start=REFERENCE_START, end=REFERENCE_END
            ), 
        color=BLUE
        )
    
    if how == 'fixed':
        plt.fill_between( 
                    ref.index,
                    ref['max_mean'], 
                    ref['max_mean'] + MIN_EXC,
                    color=LIGHTBLUE,
                    alpha = 0.2,
                    label='{}°C above average'.format(MIN_EXC)
                    )
        # Area between 5°C and 10°C above average
        #plt.fill_between( 
        #            ref.index,
        #            ref['max_mean'] + MIN_EXC, 
        #            ref['max_mean'] + 2*MIN_EXC,
        #            color=LIGHTBLUE,
        #            alpha = 0.1,
        #            label='{}°C above average'.format(2* MIN_EXC)
        #            )
        
    elif how == 'std':
        plt.fill_between( 
                    ref.index,
                    ref['max_mean'], 
                    ref['max_mean'] + ref['max_std'],
                    color=BLUE,
                    alpha = 0.2,
                    label='std above average'.format(MIN_EXC)
                    )
        # Area of 2*std
        #plt.fill_between( 
        #            ref.index,
        #            ref['max_mean'] + ref['max_std'], 
        #            ref['max_mean'] + 2*ref['max_std'],
        #            color=BLUE,
        #            alpha = 0.1,
        #            label='2*std above average'.format(MIN_EXC)
        #            )
        
        # Area of 3*std                    
        #plt.fill_between( 
        #            ref.index,
        #            ref['max_mean'] + 2*ref['max_std'], 
        #            ref['max_mean'] + 3*ref['max_std'],
        #            color=BLUE,
        #            alpha = 0.05,
        #            label='3*std above average'.format(MIN_EXC)
        #            )

    ax.set_prop_cycle(color=plt.cm.Dark2.colors)

    for year in years:
        compare_year = compare[compare['year']==year].reset_index()
        plt.plot(compare_year['heatwave_temp'], label=str(year))

    ax.set_xlim(1, 365)
    ax.set_xticks(month_starts)
    empty = '               '
    ax.set_xticklabels([empty + month for month in month_names])

    ax.yaxis.set_major_locator(ticker.MultipleLocator(5))
    ax.yaxis.set_major_formatter(
        ticker.FuncFormatter(lambda x, _: f'{x:.0f}°C')
        )

    ax.grid(True, color='gray', alpha=0.3)

    if how == 'fixed':
        method = str(MIN_EXC)+'°C'
    else:
        method = 'std'
    ax.set_title(
        'Daily maximum temperatures during heatwaves '
        'in {city} ({method})'.format(city=city, method=method),
        fontweight="bold", size=14, pad=10
        )
    plt.legend()
    plt.show()
    plt.clf()
    

def daily_max_plot(city):
    '''
    Plots the daily maximum temperatures for 2025 and 2026, 
    in comparison with the average of the reference period 
    and standard deviation.
    '''

    df, daily, ref = prepare_data(city)

    # --- PLOT ---
    fig = plt.figure(figsize=FIGSIZE)
    days = range(1,366)

    # --- X axis: ticks at 1st of each month ---
    month_starts = []
    month_labels = [''] * 13 # hide tick labels
    month_names = [
                    'Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 
                    'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec',''
                    ]

    base = dt.date(2026, 1, 1)
    for m in range(1, 13):
        d = dt.date(2026, m, 1)
        day_of_year = (d - base).days
        month_starts.append(day_of_year)

    month_starts.append(365) # Add Dec 31 as final boundary

    # --- plot ---
    ax = plt.subplot()

    plt.plot(
        ref['max_mean'], 
        label='Average of {start} – {end}'.format(
            start=REFERENCE_START, end=REFERENCE_END
            ),
        color=BLUE
        )
    plt.fill_between( 
                ref.index,
                ref['max_mean'] - ref['max_std'], 
                ref['max_mean'] + ref['max_std'],
                color=BLUE,
                alpha = 0.2,
                label='std'
                )
    #plt.fill_between( 
    #            ref.index,
    #            ref['max_mean'] - 2*ref['max_std'], 
    #            ref['max_mean'] + 2*ref['max_std'],
    #            color=BLUE,
    #            alpha = 0.1,
    #            label='2*std'
    #            )
    #plt.fill_between( 
    #            ref.index,
    #            ref['max_mean'] - 3*ref['max_std'], 
    #            ref['max_mean'] + 3*ref['max_std'],
    #            color=BLUE, 
    #            alpha = 0.05,
    #            label='3*std'
    #            )
    
    plt.plot(
        daily[daily['year']==2025].reset_index()['max'], 
        label='2025', 
        color=ORANGE
        )
    plt.plot(
        daily[daily['year']==2026].reset_index()['max'], 
        label='2026', 
        color=MAROON
        )

    ax.set_xlim(1, 365)
    ax.set_xticks(month_starts)
    empty = '               '
    ax.set_xticklabels([empty + month for month in month_names])

    ax.yaxis.set_major_locator(ticker.MultipleLocator(5))
    ax.yaxis.set_major_formatter(
        ticker.FuncFormatter(lambda x, _: f'{x:.0f}°C')
        )

    ax.grid(True, color='gray', alpha=0.3)

    ax.set_title(
        'Daily maximum temperatures in {}'.format(city),
        fontweight="bold", size=14, pad=10
        )
    plt.legend()
    plt.show()
    plt.clf()

def spaghetti_plot(city):
    '''
    Plots the Copernicus 'spaghetti' chart for daily maximum 
    temperatures and year 2026.
    '''

    df, daily, ref = prepare_data(city)

    # --- PLOT ---
    fig = plt.figure(figsize=FIGSIZE)
    days = range(1,366)

    # --- X axis: ticks at 1st of each month ---
    month_starts = []
    month_labels = [''] * 13 # hide tick labels
    month_names = [
                    'Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 
                    'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec',''
                    ]

    base = dt.date(2026, 1, 1)
    for m in range(1, 13):
        d = dt.date(2026, m, 1)
        day_of_year = (d - base).days
        month_starts.append(day_of_year)

    month_starts.append(365) # Add Dec 31 as final boundary

    # --- plot ---
    ax = plt.subplot()
    for year in range(1941, 2026):
        data_year = daily[daily['year']==year].reset_index()
        plt.plot(data_year['max'], color=GRAY, linewidth=0.5)

    plt.plot(
        ref['max_mean'], 
        label='Average of {start} – {end}'.format(
            start=REFERENCE_START, end=REFERENCE_END
            ), 
        color=BLUE, 
        linestyle='--'
        )
    data_2025 = daily[daily['year']==2025].reset_index()
    plt.plot(
        data_2025['max'], 
        color=GRAY, 
        linewidth=0.5, 
        label='1941 – 2025'
        )
        # this is a trick how to show one line legend 
        # for all the thin gray lines
    
    data_2026 = daily[daily['year']==2026].reset_index()
    plt.plot(
        data_2026['max'], 
        label='2026', 
        color=MAROON, 
        linewidth=2
        )

    ax.set_xlim(1, 365)
    ax.set_xticks(month_starts)
    empty = '              '
    ax.set_xticklabels([empty + month for month in month_names])

    ax.yaxis.set_major_locator(ticker.MultipleLocator(5))
    ax.yaxis.set_major_formatter(
        ticker.FuncFormatter(lambda x, _: f'{x:.0f}°C')
        )

    ax.grid(True, color='gray', alpha=0.3)

    ax.set_title(
        'Daily maximum temperatures in Milan',
        fontweight="bold", size=14, pad=10
        )
    plt.legend()

def prepare_data_stripes(city):
    '''
    Loads the data for the given city from "open-meteo-city.csv"
    Prepares the data for use for drawing the warming stripes
    '''
    
    # --- Read data ---
    df = pd.read_csv('open-meteo-{}.csv'.format(city), skiprows=3)

    # --- Prepare the dataset ---
    df.rename(
        columns={
            'temperature_2m (°C)' : 'temperature', 'time' : 'datetime'
            }, 
        inplace=True
        )
    df['datetime'] = pd.to_datetime(df['datetime'])
    df['year'] = df['datetime'].dt.year
    df = df[
        (df['year'] < 2026) & (df['year'] > 1940)
        ] # delete incomplete years 1940 and 2026

    # --- Compute yearly mean temperatures ---
    yearly = df.groupby(['year'])['temperature'].mean().reset_index()

    # --- Compute the reference average temperature ---
    reference = yearly[
                        (yearly['year'] >= REFERENCE_START)
                        & (yearly['year'] <= REFERENCE_END)
                        ]
    reference_average = reference['temperature'].mean()

    # --- Compute anomalies ---
    yearly['anomalies'] = yearly['temperature'] - reference_average
    
    return yearly, reference_average

def warming_stripes_plot(city):
    '''
    Plots the warming stripes for the given city
    '''

    df, avg = prepare_data_stripes(city)

    # Create the figure and axes objects, specify the size 
    # and the dots per inches 
    fig, ax = plt.subplots(figsize=(10,3), dpi = 96)

    # Colours - Choose the colour map - 8 blues and 8 reds
    cmap = ListedColormap([
        '#08306b', '#08519c', '#2171b5', '#4292c6',
        '#6baed6', '#9ecae1', '#c6dbef', '#deebf7',
        '#fee0d2', '#fcbba1', '#fc9272', '#fb6a4a',
        '#ef3b2c', '#cb181d', '#a50f15', '#67000d'])

    # linearly normalizes data into the [0.0, 1.0] interval
    norm = mpl.colors.Normalize(
            df['anomalies'].min(), 
            df['anomalies'].max()
            )

    # Plot bars
    bar = ax.bar(
        df['year'], 
        1, 
        color=cmap(norm(df['anomalies'])), 
        width=1, 
        zorder=2
        )

    # Remove the spines
    ax.spines[
        ['top', 'left', 'bottom', 'right']
        ].set_visible(False)

    # Reformat y-axis label and tick labels
    ax.set_ylabel('', fontsize=12, labelpad=10)
    ax.set_yticks([])
    ax.set_ylim([0, 1]) 

    # Adjust the margins around the plot area
    plt.subplots_adjust(
        left=0.1, 
        right=None, 
        top=None, 
        bottom=0.2, 
        wspace=None, 
        hspace=None
        )

    # Set a white background
    fig.patch.set_facecolor('white')
    ax.patch.set_facecolor('white')

    # Reformat x-axis label and tick labels
    ax.set_xlabel('', fontsize=12, labelpad=10)
    ax.xaxis.set_tick_params(
        pad=2, 
        labelbottom=True, 
        bottom=True, 
        labelsize=12, 
        labelrotation=0, 
        color='white'
        )
    ax.set_xlim([df['year'].min(), df['year'].max()+1])
    ax.xaxis.set_major_locator(ticker.MultipleLocator(10))

    # Set graph title
    ax.set_title(
        'Temperature change in {city} '
        '(1941 - 2025)'.format(city=city), 
        #loc='left', 
        fontweight="bold", size=14, pad=10
        )

    # Adjust the margins around the plot area
    plt.subplots_adjust(
        left=0.11, 
        right=None, 
        top=None, 
        bottom=0.2, 
        wspace=None, 
        hspace=None
        )