"""
linreg - GUI linear regression tools for kinetics

Copyright (C) 2024-2026 Johann Rohwer, all rights reserved

Permission to use, modify, and distribute this software is given under
the terms of the 3-clause BSD licence.
https://opensource.org/license/bsd-3-clause
"""

__doc__="""
linreg - GUI linear regression tools for kinetics
-------------------------------------------------
Analyse enzyme kinetic timecourse data with linear regression and save
the slopes (rates) as well as graphs of data.

Input:
======
`data` - list of Pandas dataframes, each with two columns,
         't' (time), and 'Abs' (absorbance)

Usage:
======
>>> from linreg import LinReg
>>> a = LinReg(data)

... use matplotlib GUI widgets for linear regressions

>>> a.make_summary()
    ... write names of individual dataframes and corresponding slopes to new dataframe

Outputs:
========
A series of PNG plots are saved in a subfolder with the name of the data list used as input.

Install requirements:
=====================
- numpy
- scipy
- pandas
- matplotlib
- ipywidgets
- ipython
"""

import numpy as np
import scipy as sp
import pandas as pd
from matplotlib import pyplot as plt
import inspect
import os

from matplotlib.widgets import SpanSelector, Button
from ipywidgets import Output
from IPython.display import display

def retrieve_name(var):
    for fi in reversed(inspect.stack()):
        names = [var_name for var_name, var_val in fi.frame.f_locals.items() if var_val is var]
        if len(names) > 0:
            for n in names:
                if n[0] != '_':
                    return n

class LinReg(object):
    i = None
    slope = None
    keep_values = {}

    def __init__(self, data):
        assert isinstance(data, list), "This class must be instantiated with a list of dataframes."
        for i in data:
            assert isinstance(i, pd.core.frame.DataFrame),  "This class must be instantiated with a list of dataframes."
        self.fig, self.ax = plt.subplots(figsize=(7, 5))
        self.ax.set(facecolor='#FFFFCC')
        self.fig.subplots_adjust(bottom=0.15, top=0.88)
        self.axprev = plt.axes([0.7, 0.02, 0.1, 0.06])
        self.axnext = plt.axes([0.81, 0.02, 0.1, 0.06])
        self.axstart = plt.axes([0.59, 0.02, 0.1, 0.06])
        self.fig.text(0.5, 0.95, 'Press left mouse button and drag to select range', ha='center')

        self.bnext = Button(self.axnext, 'Next')
        self.bnext.on_clicked(self.nextsample)
        self.bprev = Button(self.axprev, 'Prev')
        self.bprev.on_clicked(self.prevsample)
        self.bstart = Button(self.axstart, 'Start')
        self.bstart.on_clicked(self.startsample)
        self.out = Output()
        display(self.out)

        self.data = data
        self.samples = [retrieve_name(i) for i in data]
        self.checkpath()
        self.span = SpanSelector(self.ax, self.onselect, 'horizontal', useblit=True, interactive=True,
                    props=dict(alpha=0.2, facecolor='red'))

    def nextsample(self, event):
        if self.i is None:
            pass
        elif self.i == len(self.samples)-1:
            self.savefig()
            self.storedata()
        else:
            self.savefig()
            self.storedata()
            self.i += 1
            self.replot()

    def prevsample(self, event):
        if self.i == 0 or self.i is None:
            pass
        else:
            self.savefig()
            self.storedata()
            self.i -= 1
            self.replot()

    def startsample(self, event):
        if self.i is not None:
            pass
        else:
            self.i = 0
            x = self.data[self.i].t
            y = self.data[self.i].Abs
            lr = sp.stats.linregress(x, y)
            m = lr.slope
            c = lr.intercept
            r = lr.rvalue
            self.slope = m
            self.line1, = self.ax.plot(x, y, 'bd', ms=2)
            self.line2, = self.ax.plot(x, m*x+c, 'g-')
            self.text1 = self.ax.text(0.78, 0.95, 'slope=%0.5f'%m, transform=self.ax.transAxes)
            self.text2 = self.ax.text(0.78, 0.89, r'$R^2$=%0.3f'%r**2, transform=self.ax.transAxes)
            self.ax.set_title(self.samples[self.i])

    def replot(self):
        x = self.data[self.i].t
        y = self.data[self.i].Abs
        xmin, xmax = self.span.extents
        if xmin == 0 and xmax == 0:
            thisx = x
            thisy = y
        else:
            indmin, indmax = np.searchsorted(x, (xmin, xmax))
            indmax = min(len(x) - 1, indmax)
            thisx = x[indmin:indmax]
            thisy = y[indmin:indmax]
        lr = sp.stats.linregress(x=thisx, y=thisy)
        m = lr.slope
        c = lr.intercept
        r = lr.rvalue
        self.slope = m
        self.line1.set_data(x, y)
        self.line2.set_data(x, m*x+c)
        self.text1.set_text('slope=%0.5f'%m)
        self.text2.set_text(r'$R^2$=%0.3f'%r**2)
        self.ax.set_title(self.samples[self.i])
        self.ax.relim()
        self.ax.autoscale_view()
        self.fig.canvas.draw()

    def storedata(self):
        key = self.samples[self.i]
        self.keep_values[key] = self.slope

    def checkpath(self):
        dirname = retrieve_name(self.data)
        if not os.path.isdir(os.path.join(os.curdir, dirname)):
            os.makedirs(os.path.join(os.curdir, dirname))
        self.outpath = os.path.join(os.curdir, dirname)

    def savefig(self):
        outfile = self.samples[self.i]+'.png'
        self.fig.savefig(os.path.join(self.outpath, outfile), dpi=100)

    def onselect(self, xmin, xmax):
        x, y = self.line1.get_data()
        indmin, indmax = np.searchsorted(x, (xmin, xmax))
        indmax = min(len(x) - 1, indmax)
        thisx = x[indmin:indmax]
        thisy = y[indmin:indmax]
        m, c, r, p, err = sp.stats.linregress(x=thisx, y=thisy)
        self.slope = m
        self.line2.set_data(x, m*x+c)
        self.text1.set_text('slope=%0.5f'%m)
        self.text2.set_text(r'$R^2$=%0.3f'%r**2)
        self.fig.canvas.draw()

    def make_summary(self):
        return pd.DataFrame.from_dict(self.keep_values, orient='index', columns=['rate'])