"""
Plots RGA data as it's generated.
"""

import datetime, time
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import os
import json

MASS_GUESSES = { # Used in legend labels.
    2: "$H_2$",
    12: "$C$",
    16: "$O$, $CH_4$",
    17: "$HO$",
    18: "$H_2O$",
    19: "$F$",
    28: "$N_2$",
    40: "$Ar$",
    44: "$CO_2$",
    178: "C_8H_18S_2"
}

SCANS_DIR = "spoofed_rga_data" # Change to the directory that the RGA saves files into.

def get_scan_paths():
    return [SCANS_DIR + "/" + i for i in os.listdir(SCANS_DIR) if i.startswith("MassSpecData")]

def scan_stream(): # todo: replace live_scan_path with a scan path index to tidy this up
    """
    Yields lines of scan filess as they are generated by the RGA,
    from both the XML and CSV sections.
    Yields None when there are not yet more lines to yield.
    """
    scan_paths = get_scan_paths()
    while not scan_paths:
        yield None # signals "check again later"
        scan_paths = get_scan_paths()
        
    scan_paths.sort() # sorted chronologically due to how the RGA generates filenames
    live_scan_path = scan_paths[-1] # scan_paths[-1] is the most recent scan path
    file = open(scan_paths[-1])
    while True:
        
        scan_paths = get_scan_paths()
        
        while True: # yield all new lines in the current file
            cursor = file.tell()
            line = file.readline()
            if line and line.endswith("\n"):
                yield line
            else: # no more lines, or line not done being written yet
                file.seek(cursor)
                break
        
        scan_paths.sort()
        if live_scan_path != scan_paths[-1]: # current scan is no longer the most recent
            file.close()
            live_scan_path = scan_paths[scan_paths.index(live_scan_path) + 1] # this code is clunky sorry >.<"
            file = open(live_scan_path)
        else: # we're waiting for new data
            yield None # "check again later"

def generate_palette(masses):
    """
    Generates a color palette for an iterable of masses.
    Returns a dictionary of the form {mass: color}.
    """
    mass_cmap = matplotlib.cm.get_cmap("plasma")
    palette = {mass: mass_cmap(i / len(masses)) for i, mass in enumerate(masses)}
    palette[999] = "grey" # total pressure
    palette[5] = "black" # Pirani pressure?
    return palette

def mass_label(m):
    """
    Returns a label for a mass trend.
    Checks MASS_GUESSES to provide a guess for the gas species.
    """
    if m == 999:
        return "Total pressure"
    elif m == 5:
        return "Pirani pressure??"
    elif m in MASS_GUESSES:
        return "{} ({}?)".format(m, MASS_GUESSES[m])
    else:
        return str(m)

## Set up figure.
fig, (trend_ax, sweep_ax) = plt.subplots(2, 1)
fig.suptitle("Live Residual Gas Analysis", size=20, weight="bold")
fig.canvas.set_window_title("Live Residual Gas Analysis")
fig.subplots_adjust(hspace=0.4) # Adjust margin between subplots.

## Set up trend subplot.
trend_ax.set_title("Trend View")
trend_ax.set_xlabel("time")
trend_ax.set_yscale("log")
trend_ax.set_ylabel("relative pressure (Pa)") # assuming that the rga is set to Pa
trend_ax.yaxis.grid(True) # Add horizontal gridlines because Dr. Howald likes them.
trend_artists = []
mass_series = {}

## Set up sweep subplot.
sweep_ax.set_title("Sweep View")
sweep_ax.set_xlabel("mass (amu)")
sweep_ax.set_ylabel("relative pressure (Pa)") # assuming that the rga is set to Pa
sweep_ax.yaxis.grid(True) # Add horizontal gridlines because Dr. Howald likes them.
#sweep_ax.xaxis.grid(True) # Uncomment this to add vertical gridlines if you want.
sweep_ax.xaxis.set_major_locator(plt.MultipleLocator(10)) # major ticks every 10 amu
sweep_ax.xaxis.set_minor_locator(plt.MultipleLocator(1)) # minor ticks every 1 amu
sweep_ax.set_yscale("log")
sweep_pressures = {} # mass: pressure

streamer = scan_stream()

def animate():
    """
    A generator, iterated upon by FuncAnimation to animate the plot.
    Implementation note: yielding allows FuncAnimation to re-render the plot display.
    """
    last_mass = None
    current_sweep_artist = None
    config_nonce = None
    legend = fig.legend(handles=[])
    while True:
        for line in streamer:
            if line is None: # No more fresh lines to consume.
                break
            if line.startswith("2"): # It's a csv line, so it has data in it. This kludge will only work until the year 3000.
                #print(line)
                raw_t, raw_m, raw_p = [i.strip() for i in line[:-2].split(", ")]
                
                m = float(raw_m)
                if m == int(m):
                    m = int(m)
                
                t = datetime.datetime.strptime(raw_t, "%Y/%m/%d %H:%M:%S.%f")
                   
                p = float(raw_p)

                # Add new data to trend.
                if m not in mass_series:
                    mass_series[m] = [t], [p]
                else:
                    mass_series[m][0].append(t)
                    mass_series[m][1].append(p)

                # Update data in sweep.
                if m != 999: # Don't include the total pressure reading in the sweep.
                    sweep_pressures[m] = p
                
                if last_mass and m < last_mass and current_sweep_artist: # We've gone backwards, so onionskin the old data.
                    current_sweep_artist.set_color("Orchid") # Change color to distinguish from newest data.
                    current_sweep_artist = None # Don't bother trying to remove the artist later.
                    for artist in sweep_ax.lines:
                        alpha = artist.get_alpha() or 1 # alpha is None by default, in which case use 1 instead
                        if alpha < 0.1:
                            artist.remove() # artist is vanishingly transparent, so just remove it.
                        else:
                            artist.set_alpha(alpha * onion_opacity) # Onionskin old data.

                last_mass = m

        # Check configuration file. (Not very efficient, but we have to read files every frame anyway.)
        with open("chamberplot_stream_config.json") as file:
            config = json.load(file)

        if config_nonce != config["nonce"]:
            print("Updating configuration.")
            config_nonce = config["nonce"]
            interesting_masses = config["interesting_masses"]
            interesting_masses.sort()
            onion_opacity = config["onion_opacity"]
            pressure_floor = config["pressure_floor"]

            # Regenerate palette
            palette = generate_palette(interesting_masses)

            # Update legend
            legend.remove()
            legend = fig.legend(
                handles=[
                    matplotlib.patches.Patch(
                        color=palette[mass],
                        label=mass_label(mass)
                    ) for mass in interesting_masses
                ]
            )

            # Adjust lower y limit to pressure_floor.
            # Currently adjusts both subplots to use the same pressure floor.
            # This could be extended to use different pressure floors for the two subplots.
            for ax in trend_ax, sweep_ax:
                ax.set_ylim(auto=True) # let matplotlib decide the top y limit
                ax.set_ylim(bottom=pressure_floor)

        
        #### Update sweep subplot

        if current_sweep_artist is not None: # if we didn't just onionskin the current sweep artist,
            current_sweep_artist.remove() # then remove it from the plot so we can replace it with a new one.
        
        (current_sweep_artist, ) = sweep_ax.plot(
            sweep_pressures.keys(),
            sweep_pressures.values(),
            color="DarkOrchid"
        )

        # Remove all markers.
        for artist in sweep_ax.lines:
            if artist.get_marker() != "None":
                artist.remove()
        
        # Mark interesting masses on the sweep subplot.
        for m in interesting_masses:
            if m not in sweep_pressures: # No sweep data for this mass, so skip it.
                continue
            
            sweep_ax.plot(
                [m],
                [sweep_pressures[m]],
                marker=matplotlib.markers.CARETDOWN,
                markersize=12,
                color=palette[m] # Use the same color as the trend does for this mass.
            )
        
        # Mark the current sweep position on the sweep subplot so it's easier to follow with the eye.
        if last_mass is not None:
            sweep_ax.plot(
                [last_mass],
                [sweep_pressures[last_mass]],
                marker="*",
                markersize=12,
                color="Indigo"
            )

        
        #### Update trend subplot

        for artist in trend_artists:
            artist.remove()
        trend_artists.clear()
        
        for m in interesting_masses:
            if m not in mass_series:
                continue # No data yet for this mass, so don't bother trying to plot it.
            
            times, pressures = mass_series[m]
            trend_artists.append(*trend_ax.plot(
                times,
                pressures,
                color=palette.get(m)
            ))
        
        # Let FuncAnimation update the display.
        yield
    

animator = animate()
ani = FuncAnimation(fig, lambda x: next(animator), 1000)
plt.show()


