import click
import numpy as np
import pandas as pd
from scipy.integrate import simps

from logs import setup_logging
from utils import crossings

logger = setup_logging()


def return_first_or_nan(l):
    """
    Returns the first value of a list if empty or returns nan. Used for getting dune/toe and crest values.
    :param l:
    :return:
    """
    if len(l) == 0:
        return np.nan
    else:
        return l[0]


def volume_change(df_profiles, df_profile_features, zone):
    """
    Calculates how much the volume change there is between prestrom and post storm profiles.
    :param df_profiles:
    :param df_profile_features:
    :param zone: Either 'swash' or 'dune_face'
    :return:
    """
    logger.info("Calculating change in beach volume in {} zone".format(zone))

    df_vol_changes = pd.DataFrame(
        index=df_profile_features.index.get_level_values("site_id").unique()
    )
    df_profiles = df_profiles.sort_index()
    sites = df_profiles.groupby(level=["site_id"])

    for site_id, df_site in sites:
        logger.debug(
            "Calculating change in beach volume at {} in {} zone".format(site_id, zone)
        )

        prestorm_row = df_profile_features.loc[(site_id, "prestorm")]
        prestorm_dune_toe_x = prestorm_row.dune_toe_x
        prestorm_dune_crest_x = prestorm_row.dune_crest_x

        # If no dune toe has been defined, Dlow = Dhigh. Refer to Sallenger (2000).
        if np.isnan(prestorm_dune_toe_x):
            prestorm_dune_toe_x = prestorm_dune_crest_x

        # If no prestorm and poststorm profiles, skip site and continue
        profile_lengths = [
            len(df_site.xs(x, level="profile_type")) for x in ["prestorm", "poststorm"]
        ]
        if any([length == 0 for length in profile_lengths]):
            continue

        # Find last x coordinate where we have both prestorm and poststorm measurements. If we don't do this,
        # the prestorm and poststorm values are going to be calculated over different lengths.
        df_zone = df_site.dropna(subset=["z"])
        x_last_obs = min(
            [
                max(
                    df_zone.query(
                        "profile_type == '{}'".format(profile_type)
                    ).index.get_level_values("x")
                )
                for profile_type in ["prestorm", "poststorm"]
            ]
        )
        x_first_obs = max(
            [
                min(
                    df_zone.query(
                        "profile_type == '{}'".format(profile_type)
                    ).index.get_level_values("x")
                )
                for profile_type in ["prestorm", "poststorm"]
            ]
        )

        # Where we want to measure pre and post storm volume is dependant on the zone selected
        if zone == "swash":
            x_min = max(prestorm_dune_toe_x, x_first_obs)
            x_max = x_last_obs
        elif zone == "dune_face":
            x_min = max(prestorm_dune_crest_x, x_first_obs)
            x_max = min(prestorm_dune_toe_x, x_last_obs)
        else:
            logger.warning("Zone argument not properly specified. Please check")
            x_min = None
            x_max = None

        # Now, compute the volume of sand between the x-coordinates prestorm_dune_toe_x and x_swash_last for both prestorm
        # and post storm profiles.
        prestorm_vol = beach_volume(
            x=df_zone.query("profile_type=='prestorm'").index.get_level_values("x"),
            z=df_zone.query("profile_type=='prestorm'").z,
            x_min=x_min,
            x_max=x_max,
        )
        poststorm_vol = beach_volume(
            x=df_zone.query("profile_type=='poststorm'").index.get_level_values("x"),
            z=df_zone.query("profile_type=='poststorm'").z,
            x_min=x_min,
            x_max=x_max,
        )

        # Identify the x location where our pre and post storm profiles start to differ. This is so changes no due to
        # the storm are not included when calculating volume.
        df_prestorm = (
            df_site.xs("prestorm", level="profile_type").z.rename("z_pre").to_frame()
        )
        df_poststorm = (
            df_site.xs("poststorm", level="profile_type").z.rename("z_post").to_frame()
        )
        df_diff = df_prestorm.merge(df_poststorm, on=["site_id", "x"])
        df_diff["z_diff"] = df_diff.z_pre - df_diff.z_post

        # Find all locations where the difference in pre and post storm is zero. Take the most seaward location as the
        # x location where our profiles are the same.
        x_crossings = crossings(df_diff.index.get_level_values("x"), df_diff.z_diff, 0)

        if len(x_crossings) != 0:
            x_change_point = x_crossings[-1]
        else:
            x_change_point = np.nan

        # # For debugging
        # import matplotlib.pyplot as plt
        # f,(ax1,ax2) = plt.subplots(2,1,sharex=True)
        # ax1.plot(df_prestorm.index.get_level_values('x'), df_prestorm.z_pre,label='prestorm')
        # ax1.plot(df_poststorm.index.get_level_values('x'), df_poststorm.z_post,label='poststorm')
        # ax1.axvline(x_crossings[-1], color='red', linestyle='--', linewidth=0.5, label='Change point')
        # ax1.legend()
        # ax1.set_title(site_id)
        # ax1.set_ylabel('elevation (m AHD)')
        # ax2.plot(df_diff.index.get_level_values('x'), df_diff.z_diff)
        # ax2.set_xlabel('x coordinate (m)')
        # ax2.set_ylabel('elevation diff (m)')
        # ax2.axvline(x_crossings[-1],color='red',linestyle='--',linewidth=0.5)
        # plt.show()

        diff_vol = beach_volume(
            x=df_diff.index.get_level_values("x"),
            z=df_diff.z_diff,
            x_min=np.nanmax([x_min, x_change_point]),
            x_max=np.nanmax([x_max, x_change_point]),
        )

        # Here, if cannot calculate the difference volume, assume no volume change
        if np.isnan(diff_vol):
            diff_vol = 0

        df_vol_changes.loc[site_id, "prestorm_{}_vol".format(zone)] = prestorm_vol
        df_vol_changes.loc[site_id, "poststorm_{}_vol".format(zone)] = poststorm_vol
        df_vol_changes.loc[site_id, "{}_vol_change".format(zone)] = diff_vol
        df_vol_changes.loc[site_id, "{}_pct_change".format(zone)] = (
            diff_vol / prestorm_vol * 100
        )

    return df_vol_changes


def beach_volume(x, z, x_min=np.NINF, x_max=np.inf):
    """
    Returns the beach volume of a profile, calculated with Simpsons rule
    :param x: x-coordinates of beach profile
    :param z: z-coordinates of beach profile
    :param x_min: Minimum x-coordinate to consider when calculating volume
    :param x_max: Maximum x-coordinate to consider when calculating volume
    :return:
    """
    profile_mask = [True if x_min < x_coord < x_max else False for x_coord in x]
    x_masked = np.array(x)[profile_mask]
    z_masked = np.array(z)[profile_mask]

    if len(x_masked) == 0 or len(z_masked) == 0:
        return np.nan
    else:
        return simps(z_masked, x_masked)


def storm_regime(df_observed_impacts):
    """
    Returns the dataframe with an additional column of storm impacts based on the Storm Impact Scale. Refer to
    Sallenger (2000) for details.
    :param df_observed_impacts:
    :return:
    """
    logger.info("Getting observed storm regimes")

    swash = (df_observed_impacts.dune_face_pct_change <= 2) & (
        df_observed_impacts.dune_face_vol_change <= 3
    )
    collision = (df_observed_impacts.dune_face_pct_change >= 2) | (
        df_observed_impacts.dune_face_vol_change > 3
    )

    df_observed_impacts.loc[swash, "storm_regime"] = "swash"
    df_observed_impacts.loc[collision, "storm_regime"] = "collision"

    # TODO We may be able to identify observed regimes by looking at the change in crest and toe elevation. This would be useful for
    # locations where we have overwash and cannot calculate the change in volume correctly. Otherwise, maybe it's better to put it in manually.

    return df_observed_impacts


def overwrite_impacts(df_observed_impacts, df_raw_features):
    """
    Overwrites calculated impacts with impacts manually specified in profile_features file
    :param df_raw_profile_features:
    :return:
    """
    df_observed_impacts.update(
        df_raw_features.rename(columns={"observed_storm_regime": "storm_regime"})
    )
    return df_observed_impacts


@click.command()
@click.option("--profiles-csv", required=True, help="")
@click.option("--profile-features-crest-toes-csv", required=True, help="")
@click.option("--raw-profile-features-csv", required=True, help="")
@click.option("--output-file", required=True, help="")
def create_observed_impacts(
    profiles_csv, profile_features_crest_toes_csv, raw_profile_features_csv, output_file
):

    profiles_csv = "./data/interim/profiles.csv"
    profile_features_crest_toes_csv = "./data/interim/profile_features_crest_toes.csv"
    raw_profile_features_csv = (
        "./data/raw/profile_features_chris_leaman/profile_features_chris_leaman.csv"
    )

    logger.info("Creating observed wave impacts")
    logger.info("Importing data")
    df_profiles = pd.read_csv(profiles_csv, index_col=[0, 1, 2])
    df_profile_features = pd.read_csv(profile_features_crest_toes_csv, index_col=[0, 1])

    logger.info("Creating new dataframe for observed impacts")
    df_observed_impacts = pd.DataFrame(
        index=df_profile_features.index.get_level_values("site_id").unique()
    )

    logger.info("Getting pre/post storm volumes")
    df_swash_vol_changes = volume_change(df_profiles, df_profile_features, zone="swash")
    df_dune_face_vol_changes = volume_change(
        df_profiles, df_profile_features, zone="dune_face"
    )
    df_observed_impacts = df_observed_impacts.join(
        [df_swash_vol_changes, df_dune_face_vol_changes]
    )

    # Classify regime based on volume changes
    df_observed_impacts = storm_regime(df_observed_impacts)

    # Overwrite storm impacts with manually picked impacts
    df_raw_features = pd.read_csv(raw_profile_features_csv, index_col=[0])
    df_observed_impacts = overwrite_impacts(df_observed_impacts, df_raw_features)

    # Save dataframe to csv
    df_observed_impacts.to_csv(output_file, float_format="%.4f")

    logger.info("Saved to %s", output_file)
    logger.info("Done!")