#!/usr/bin/env python3
import plotly.express as px
from pandas import DataFrame

# df = pd.DataFrame(data=dict(x=[1, 2, 3, 4], y=[1, 4, 9, 16]))
# fig = px.line(df, x="x", y="y", title="test")
# fig.show()

class Line:
    def __init__(self, from_, to):
        self.from_ = from_
        self.to = to

    def slope(self) -> float:
        return (self.to[1] - self.from_[1]) / (self.to[0] - self.from_[0])

    def get(self, x: float) -> float:
        from_x, from_y = self.from_
        to_x, to_y = self.to
        tween = (x - from_x) / (to_x - from_x)
        return tween * to_y + (1-tween) * from_y


class Piecewise:
    def __init__(self, xy: list):
        """ xy is a list of (x, y) pairs """
        self.xy = list(xy)

    def normalized(self, prev_max: float) -> 'Piecewise':
        """ map every coordinate from [-prev_max, prev_max] to [0, 1] """
        p = prev_max
        r = 2*prev_max
        s = 1.0/r
        return Piecewise([
            (s*(x + p), s*(y + p)) for (x, y) in self.xy
        ])

    def logically_inverted(self) -> 'Piecewise':
        """ return a Piecewise that evaluates to 1-y """
        return Piecewise([
            (x, 1-y) for (x, y) in self.xy
        ])

    def line_for(self, x: float) -> Line:
        for first_lower in self.xy[:-1][::-1]:
            if first_lower[0] < x: break
        for first_upper in self.xy[1:]:
            if first_upper[0] > x: break

        return Line(first_lower, first_upper)

    def get(self, x: float) -> float:
        """
        evaluate the piecewise function at the provided x value.
        OOB points are just extrapolated from the nearest piece.
        """
        return self.line_for(x).get(x)

    def get_slope(self, x: float) -> float:
        return self.line_for(x).slope()

    def get_repeated(self, x: float, n: int = 63) -> float:
        for _ in range(n):
            x = self.get(x)
        return x

    def df_for(self, from_: float, to: float, points: int, f) -> DataFrame:
        x_step = (to - from_) / (points - 1)
        x = [from_ + x_step*x for x in range(points)]
        y = [f(xi) for xi in x]
        return DataFrame(data=dict(x=x, y=y))

    def df(self, from_: float = 0.0, to: float = 1.0, points: int = 101) -> DataFrame:
        return self.df_for(from_, to, points, self.get)

    def slope_df(self, from_: float = 0.0, to: float = 1.0, points: int = 101) -> DataFrame:
        return self.df_for(from_, to, points, self.get_slope)

    def plot_for(self, from_: float, to: float, title: str, f):
        df = self.df_for(from_, to, points=101, f=f)
        fig = px.line(df, x="x", y="y", title=title)
        fig.show()

    def plot(self, from_: float = 0.0, to: float = 1.0, title: str = "Piecewise"):
        self.plot_for(from_, to, title, self.get)

    def plot_slope(self, from_: float = 0.0, to: float = 1.0, title: str = "Piecewise"):
        self.plot_for(from_, to, title, self.get_slope)

    def plot_equilibrium(self, from_: float = 0.0, to: float = 1.0, title: str = "Piecewise"):
        self.plot_for(from_, to, title, self.get_repeated)


fwd_26 = Piecewise(
    [
        [ -14687, -7326 ],
        [ -13049, -6503 ],
        [ -11785, -5833 ],
        [ -4649,  -1447 ],
        [ 4961,    7059 ],
        [ 11283,  11147 ],
    ]
).normalized(17000)

# fwd_26.plot(title = "26 forward")
# fwd_26.logically_inverted().plot(title = "26 inverted")
# fwd_26.plot_slope(title = "26 slope")
# fwd_26.logically_inverted().plot_equilibrium(title = "26 equilibrium")


fwd_38_2_0 = Piecewise(
    [
        [ (-13745 + -13012)/2, -6222 ],
        [ (-4969 + -4744)/2,    2373 ],
        [ (1772 + 2070)/2,     10467 ],
        [ (4472 + 4114)/2,     12921 ],
        [ (7221 + 6291)/2,     14530 ],
        [ (11159 + 10397)/2,   15865 ],
        [ (12430 + 15653)/2,   16202 ],
    ]
).normalized(17000)
fwd_38_2_0.plot(title = "38 2:0 forward")
fwd_38_2_0.logically_inverted().plot(title = "38 2:0 inverted")
fwd_38_2_0.plot_slope(title = "38 2:0 slope")
fwd_38_2_0.logically_inverted().plot_equilibrium(title = "38 2:0 equilibrium")

fwd_38_3_0 = Piecewise(
    [
        [ (-13956 + -13890 + -13077)/3, -5203],
        [ (-4979 + -4885  + -4717)/3, 5051],
        [ (1531 + 503 + 1006)/3, 12509],
        [ (4180 + 1821 + 2239)/3, 14386],
        [ (6986 + 3436 + 3701)/3, 15451],
        [ (10482 + 6644 + 7735)/3, 16081],
        [ (11436 + 13343 + 14411)/3, 16380],
    ]
).normalized(17000)
fwd_38_3_0.plot(title = "38 3:0 forward")
fwd_38_3_0.logically_inverted().plot(title = "38 3:0 inverted")
fwd_38_3_0.plot_slope(title = "38 3:0 slope")
fwd_38_3_0.logically_inverted().plot_equilibrium(title = "38 3:0 equilibrium")