import sys

import matplotlib.pylab as plt
import seaborn as sns
from wela.dataloader import dataloader
from wela.imageviewer import ImageViewer, get_h5files
from wela.plotting import kymograph
from wela.sorting import sort_by_budding

h5dir = "/Users/pswain/wip/aliby_output/ribosomes/"
omids = [
    "1999_2024_01_14_Rpl3_glc_raf_00",
    "2003_2024_01_29_2pc_switch_00",
    "2008_2024_02_12_Rpl3_glc_raf_switching_00",
]
# FILL IN
server_info = {
    "host": "staffa.bio.ed.ac.uk",
    "username": "",
    "password": "",
}
view = True
# pick the experiment to analyse
omid = 2008

# 1. Run with view=True to check visually that aliby has worked correctly.
# 2. Set key_index, the signal you are most interested in.
# 3. Run with view=False to run dataloader and save a tsv file.

if view:
    omero_name = [om for om in omids if str(omid) in om][0]
    h5files = get_h5files(h5dir, omero_name)
    position = h5files[0]
    h5file = f"{h5dir}{omero_name}/{position}"
    iv = ImageViewer.remote(h5file, server_info, omid)
    tpt_end = 10
    no_cells = 6
    iv.view(
        trap_ids=iv.sample_traps_with_cells(
            tpt_end=tpt_end, no_cells=no_cells
        ),
        tpt_end=tpt_end,
        channels_to_skip=["cy5"],
        no_rows=2,
    )
    sys.exit(0)
else:
    # run dataloader
    key_index = "median_GFP"
    dl = dataloader(h5dir, ".")
    expt = [omid_full for omid_full in omids if str(omid) in omid_full][0]
    dl.load(expt, key_index=key_index, cutoff=0.9)
    dl.save()

    # plot kymographs
    groups = dl.df.group.unique()
    for group in groups:
        _, buddings = dl.get_time_series("buddings", group=group)
        sort_order = sort_by_budding(buddings)
        kymograph(
            dl.df[dl.df.group == group],
            hue=key_index,
            title=group,
            sort_order=sort_order,
        )

    # plot means
    sns.relplot(data=dl.df, x="time", y=key_index, kind="line", hue="group")
    plt.show()