import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from pandas import Timedelta
from pathlib import Path

from epilepsy_bids_loader import ManifestBIDSLoader


def as_hours(time_delta: Timedelta):
    return time_delta.total_seconds() / 3600

def in_secs(duration: float):
    return Timedelta(duration, unit="sec")

def plot_events(
    loader: ManifestBIDSLoader,
    save_path: Path,
    title: str,
):
    bids_tree = loader.data
    sub_ses = [
        (sub, ses)
        for sub, sub_data in bids_tree.items()
        for ses in sub_data.keys()
    ]

    fig_width = 64
    fig_height = len(sub_ses) * 2

    fig = plt.figure(figsize=(fig_width, fig_height))
    gs = GridSpec(nrows=len(sub_ses), ncols=1)
    axes = []
    ds_durations = []
    ds_sz_durations = []

    for i, (sub, ses) in enumerate(sub_ses):
        ax = fig.add_subplot(gs[i, 0])
        ses_data = bids_tree[sub][ses]
        ses_datetime = None
        ses_durations = []
        ses_sz_durations = []
        for run in ses_data.keys():
            events = loader.events.loc[[(sub, ses, run)]]
            for _, event in events.iterrows():
                if ses_datetime is None:
                    ses_datetime = event["dateTime"]
                start = event["dateTime"]
                end = (
                    event["dateTime"]
                    + in_secs(event["recordingDuration"])
                )
                event_type = event["eventType"]
                xmin = as_hours(start - ses_datetime)
                xmax = as_hours(end - ses_datetime)
                ds_durations.append(event["recordingDuration"])
                ses_durations.append(event["recordingDuration"])

                ax.plot(
                    [xmin, xmax],
                    [0, 0],
                    c="green" if event_type == "bckg" else "black",
                    lw=5
                )
                ax.text(
                    x=xmin,
                    y=-0.9,
                    s=run,
                    rotation=90,
                    ha="left",
                    va="center"
                )
                ax.axvline(
                    x=xmin,
                    ls="dotted",
                    c="dimgray"
                )

                if event["eventType"] == "bckg":
                    continue

                onset = start + in_secs(event["onset"])
                offset = onset + in_secs(event["duration"])
                sz_xmin = as_hours(onset - ses_datetime)
                sz_xmax = as_hours(offset - ses_datetime)
                ds_sz_durations.append(event["duration"])
                ses_sz_durations.append(event["duration"])
                ax.axvspan(
                    xmin=sz_xmin,
                    xmax=sz_xmax,
                    color="red"
                )
                ax.text(
                    x=sz_xmin,
                    y=0.9,
                    s=event_type,
                    rotation=90,
                    ha="right",
                    va="center"
                )

                continue
            continue

        ax.set_ylim([-1.8, 1.8])
        ax.set_yticks([0])
        ax.set_xlim([0, xmax])
        ses_sz_ratio = sum(ses_sz_durations) / sum(ses_durations) * 100
        ax.set_yticklabels([f"{sub},{ses}\n{ses_sz_ratio:.2f}%"])
        axes.append(ax)
        continue

    sz_dur = sum(ds_sz_durations) / 3600
    total_dur = sum(ds_durations) / 3600
    sz_ratio = sz_dur / total_dur * 100
    axes[0].set_title(
        title
        + f", total_dur = {total_dur:.1f} hrs"
        + f", sz_dur = {sz_dur:.1f} hrs ({sz_ratio:.2f}%)"
    )
    plt.tight_layout()
    plt.savefig(save_path)

    return


def main(args):
    loader = ManifestBIDSLoader(
        data_path=args.path,
        manifest_path=args.manifest,
        enforce_chronology=args.chrono,
        allowed_subjects=None
    )
    plot_events(
        loader=loader,
        title=args.path,
        save_path=args.outpath
    )
    return


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-p", "--path", type=str,
        # required=True,
        default="data/BIDS_Siena"
        # default="data/BIDS_CHB-MIT"
    )
    parser.add_argument(
        "-m", "--manifest", type=str,
        # required=True,
        default="manifests/siena.yaml"
        # default="manifests/chbmit.yaml"
    )
    parser.add_argument(
        "-c", "--chrono", type=bool, nargs='?', const=True,
        default=True,
        # default=False,
    )
    parser.add_argument(
        "-o", "--outpath", type=str,
        # required=True,
        default="tests/siena_events.pdf"
        # default="tests/chbmit_events.pdf"
    )
    args = parser.parse_args()
    main(args)