# importing packages
import pandas as pd
import numpy as np

# getting the stations 
url = f"https://raw.githubusercontent.com/PhilChodrow/PIC16B/master/datasets/noaa-ghcn/station-metadata.csv"
stations = pd.read_csv(url)

# getting the country codes data
url = f"https://raw.githubusercontent.com/mysociety/gaze/master/data/fips-10-4-to-iso-country-codes.csv"
countries = pd.read_csv(url)

# getting the temperature data
temps = pd.read_csv("/Users/shreeshagarwal/Desktop/PIC16/Datasets/temps.csv")

def prepare_df(df):
    df = df.set_index(keys=["ID", "Year"])
    df = df.stack()
    df = df.reset_index()
    df = df.rename(columns = {"level_2"  : "Month" , 0 : "Temp"})
    df["Month"] = df["Month"].str[5:].astype(int)
    df["Temp"]  = df["Temp"] / 100
    return(df)

temps = prepare_df(temps)
temps.head()

import sqlite3 

# create a new database
conn = sqlite3.connect("temps.db")

# adding the three dataframes to the database

# adding temps.csv
df_iter = pd.read_csv("/Users/shreeshagarwal/Desktop/PIC16/Datasets/temps.csv", chunksize = 100000)
for df in df_iter:
    df = prepare_df(df)
    df.to_sql("temperatures", conn, if_exists = "append", index = False)
    
# adding the stations data frame
stations.to_sql("stations", conn, if_exists = "replace", index = False)

# adding the countries data frame
countries.to_sql("countries", conn, if_exists = "replace", index = False)

def query_climate_database(country, year_begin, year_end, month):
    
    cmd = \
    f"""
    SELECT S.name, S.latitude, S.longitude, C.name, T.year, T.month, T.temp
    FROM temperatures T 
    LEFT JOIN stations S ON T.id = S.id
    LEFT JOIN countries C ON SUBSTRING(T.id,1, 2) = C.`FIPS 10-4` 
    WHERE T.year >= {year_begin} AND T.year <= {year_end} AND T.month == {month} AND C.name == "{country}"
    """
    return pd.read_sql_query(cmd, conn)

india_temps = query_climate_database("India", 1980, 2020, 1)
india_temps.head()

import sklearn
from sklearn.linear_model import LinearRegression

def coef(data_group):
  """
  This function calculates a Linear Regression coefficient 
  given the predictor variable "Year" and response variable "Temps"
  """

    x = data_group[["Year"]] # 2 brackets because X should be a df
    y = data_group["Temp"]   # 1 bracket because y should be a series
    LR = LinearRegression()
    LR.fit(x, y)
    return LR.coef_[0]

def temperature_coefficient_plot(country, year_begin, year_end, month, min_obs, **kwargs):

  """
  This function creates a 'carto-positron' style map which shows 
  stations in the specified country with their average 
  annual increase in temperature coefficient 
  measured between year_begin and year_end 
  for the specified month.
  """
    
    # SQL query to get the daataset with country, year_begin, year_end, month specifications
    temps = query_climate_database(country, year_begin, year_end, month)
    
    # removing the stations with less than min_obs
    df = temps.groupby(["NAME"])[["Temp"]].aggregate(len) 
    names = (df[df["Temp"] > 30].reset_index())["NAME"]
    
    temps = temps[temps["NAME"].isin(np.array(names))]
    
    # creating linear coefficients 
    coefs = temps.groupby("NAME").apply(coef)
    coefs = coefs.reset_index()
    
    # merge coefs with latitude and longitude columns 
    df = pd.merge(coefs, temps, on = ["NAME"])
    df[0] = df[0].round(5)
    df = df.rename(columns = {0  : "Estimated Yearly Increase in Temperature (°C)"}) 

    # drop extra columns 
    df = df.drop(["Name","Year","Temp","Month"], axis = 1)  

    # drop the duplicate entries
    df = df.drop_duplicates(subset=['NAME']) 
    
    # creating the figure
    color_map = px.colors.diverging.RdGy_r
    
    month_dict = { 1: "January", 2: "February", 3: "March", 4: "April", 5: "May", 6: "June", 7: "July", 8: "August", 9: "September", 10: "October", 11: "November", 12: "December"}
    
    fig = px.scatter_mapbox(df,
                            lat = "LATITUDE",
                            lon = "LONGITUDE",
                            hover_name = "NAME",
                            color = "Estimated Yearly Increase in Temperature (°C)",
                            mapbox_style = "carto-positron",
                            color_continuous_scale = color_map,
                            zoom = 2,
                            title = f"Estimated Yearly Increase in Temperature (°C) in {month_dict[1]} for stations in {country}",
                            **kwargs
                            )
    
    write_html(fig, "Fig1.html")
    
    return fig

temperature_coefficient_plot("India", 1980, 2020, 1, 30)

def plot2_sql_query(country, year_begin, year_end):
    
    cmd = \
    f"""
    SELECT C.name, T.year, T.month, T.temp
    FROM temperatures T 
    LEFT JOIN stations S ON T.id = S.id
    LEFT JOIN countries C ON SUBSTRING(T.id,1, 2) = C.`FIPS 10-4` 
    WHERE T.year >= {year_begin} AND T.year <= {year_end} AND C.name == "{country}"
    """
    
    return pd.read_sql_query(cmd, conn)

"""
This function calculates a z-score given an array x
"""
def z_score_calc(x):
    m = np.mean(x)
    s = np.std(x)
    return (x - m)/s

def plot_heatmap(country, year_begin, year_end, z_score = False, col_scheme = 'icefire'):

  """
  Creates a heatmap visualizing the average temperature measurements 
  for a country between year_begin and year_end for each month.
  """
    
    # get the dataframe
    df = plot2_sql_query(country, 2000,2020)
    
    # calculate average temperature by year and month 
    df = df.sort_values(by=['Year', 'Month'])
    df = df.reset_index()
    df = df.drop(["index"], axis = 1)
    
    # PLOT THE Z-SCORE HEATMAP 
    if z_score == True:
        
        # calculate z_scores and add the column to the df
        df["z"] = df.groupby(["Year","Month"])["Temp"].transform(z_score_calc)
        df = df.groupby(["Year","Month"])["z"].aggregate(np.mean).to_frame()
        df = df.reset_index()
        
        # make the final df for plotting
        plot_df = df.pivot("Year","Month","z")

        # plot the heatmap
        fig = px.imshow(plot_df, 
                    labels = dict(x = "Month", y = "Year", color = "Z score"),
                    color_continuous_scale = col_scheme,
                    x = months,
                    y = np.arange(year_begin, year_end + 1),
                    title = f"Heatmap showing Temperature (°C) z-scores of {country} between {year_begin} and {year_end} for each Month"
                    )
        
        # change the plot size
        fig.update_layout(height = 1000, width = 1000)
        
        
        return fig

    # PLOT THE AVG TEMPERATURE HEATMAP
    
    # group by year and month and calculate the mean
    df = df.groupby(["Year","Month"]).aggregate(np.mean)
    df = df.reset_index()
    
    # make the final df for plotting
    plot_df = df.pivot("Year","Month","Temp")
    
    # plot the heatmap
     fig = px.imshow(plot_df, 
                    labels = dict(x = "Month", y = "Year", color = "Temperature"),
                    color_continuous_scale = col_scheme,
                    x = months,
                    y = np.arange(year_begin, year_end + 1),
                    title = f"Heatmap showing Temperature (°C) of {country} between {year_begin} and {year_end} for each Month"
                    )
    
    # change the plot size
    fig.update_layout(height = 1000, width = 1000)
    
    
    return fig

plot_heatmap("Iceland", 2000,2020)

plot_heatmap("Iceland", 2000,2020, z_score = True)

# plot the normal heatmap
    fig, ax = plt.subplots(figsize= (17,13))
    ax.set_title(f"Heatmap showing Average Temperature (°C) of {country} between {year_begin} and {year_end} for each Month")
    fig = sns.heatmap(plot_df, ax = ax, linewidths=.5, cmap= "coolwarm", annot = True)

# plot the z-score heatmap
    fig, ax = plt.subplots(figsize= (17,13))
        ax.set_title(f"Heatmap showing Temperature (°C) z-scores of {country} between {year_begin} and {year_end} for each Month")
        fig = sns.heatmap(plot_df, ax = ax, linewidths=.5, cmap = "coolwarm")

def plot3_sql_query(country, decade_start, decade_end):
    
    """
    Performs a SQL query to obtain temperature data of 
    the specified country between specified dates.
    
    Note: Input to decade_start, decade_end should be the last year of the decade eg 2019.
    
    """
    
    cmd = \
    f"""
    SELECT C.name, T.year, T.temp
    FROM temperatures T 
    LEFT JOIN stations S ON T.id = S.id
    LEFT JOIN countries C ON SUBSTRING(T.id,1, 2) = C.`FIPS 10-4` 
    WHERE T.year >= {decade_start} AND T.year <= {decade_end} AND C.name == "{country}"
    """
    
    return pd.read_sql_query(cmd, conn)

def plot_boxplots(country, decade_begin, decade_end):

  """
  This function creates a boxplot of temperatures of a country
  for every decade between decade_begin and decade_end.
  """
    
    # getting the data 
    df = plot3_sql_query(country, decade_begin, decade_end)
    
    # making the decades array
    years = np.arange(decade_begin, decade_end + 1).tolist()
    decades = [years[i:i+10] for i in range(0, len(years), 10)]
    
    # empty dataframe for stacking 
    df2 = pd.DataFrame()

    # making and appending the decade column
    for i in range(0,len(decades)):
        df_d1 = df[df["Year"].isin(decades[i])]
        df_d1["Decade"] = f"{decades[i][0]} - {decades[i][9]}"
        frames = [df2, df_d1]
        df2 = pd.concat(frames)
        
    # create the plot
    fig = px.box(df2, 
                 y = "Decade", 
                 x = "Temp",
                 title= f"Boxplot showing Median Temperature (°C) of {country} 
                 for each decade between {decade_begin} and {decade_end + 1}",)
    
    # save the plot 
    write_html(fig, "Fig4.html")
    
    return fig

plot_boxplots("Iceland", 1970, 2019)