from typing import TextIO, Optional, Dict
from numbers import Number

import numpy as np
from numpy.typing import ArrayLike

import pandas as pd
from curvaceous.interval import Interval

Ys = Dict[str, np.ndarray]


class Curve:
    xs: np.ndarray
    ys: Ys

    def __init__(self, xs, ys: Ys):
        if not isinstance(xs, np.ndarray):
            xs = np.array(xs)

        new_ys = {}
        for k, v in ys.items():
            if not isinstance(v, np.ndarray):
                v = np.array(v)
            if len(xs) != len(v):
                raise ValueError("Coordinate lists need to be of equal length")
            new_ys[k] = v

        if not new_ys:
            raise ValueError("Curve must contain at least one y series")

        self.xs = xs
        self.ys = new_ys

    def __len__(self):
        return len(self.xs)

    @classmethod
    def from_dict(cls, m):
        return cls(m["xs"], {k: np.array(v) for k, v in m["ys"].items()})

    @property
    def domain(self):
        return Interval(np.min(self.xs), np.max(self.xs))

    def __call__(self, c, key: str, *, cut=False) -> Optional[float]:
        ys = self.ys[key]

        if cut:
            c = self.domain.closest(c)
        if c not in self.domain:
            return None

        i = np.argmax(self.xs > c)

        if i == 0:
            return ys[-1]

        width = self.xs[i] - self.xs[i - 1]
        u = (c - self.xs[i - 1]) / width
        v = (self.xs[i] - c) / width

        return v * ys[i - 1] + u * ys[i]

    def __rmul__(self, other):
        if not isinstance(other, (int, float)):
            raise TypeError("Only scalar multiplication is supported")
        ys = {k: other * v for k, v in self.ys.items()}
        return Curve(self.xs, ys)

    def scale(self, factor, key):
        if not isinstance(factor, (int, float)):
            raise TypeError("Can only scale by a scalar")
        return Curve(self.xs / factor, {key: self.ys[key]})

    def __add__(self, other):
        if isinstance(other, Curve):
            xs = np.unique(np.concatenate((self.xs, other.xs)))
            ys = {}

            for key in self.ys:
                ys[key] = np.array([
                    self(x, key) + other(x, key) for x in xs
                ])

            return Curve(xs, ys)

        if isinstance(other, (int, float)):
            ys = {k: v + other for k, v in self.ys.items()}
            return Curve(self.xs, ys)

        raise ValueError()

    def __sub__(self, other):
        if isinstance(other, Curve):
            xs = np.unique(np.concatenate((self.xs, other.xs)))
            ys = {}

            for key in self.ys:
                ys[key] = np.array([
                    self(x, key) - other(x, key) for x in xs
                ])

            return Curve(xs, ys)

        if isinstance(other, (int, float)):
            ys = {k: v - other for k, v in self.ys.items()}
            return Curve(self.xs, ys)

        raise ValueError()

    def __eq__(self, other):
        if not isinstance(other, Curve):
            return False

        if not np.array_equal(self.xs, other.xs):
            return False

        if set(self.ys.keys()) != set(other.ys.keys()):
            return False

        return all(np.array_equal(self.ys[k], other.ys[k]) for k in self.ys)

    def resample(self, cs, key):
        xs = []
        ys = []

        for x in cs:
            if x in self.domain:
                xs.append(x)
                ys.append(self(x, key))

        return Curve(xs, {key: np.array(ys)})

    def derivative(self, key):
        y = self.ys[key]
        hs = y[1:] - y[:-1]
        ws = self.xs[1:] - self.xs[:-1]
        return Curve(self.xs[:-1], {key: hs / ws})

    def grow(self, factor):
        ys = {k: factor * v for k, v in self.ys.items()}
        return Curve(factor * self.xs, ys)

    def to_dict(self):
        return {
            "xs": list(self.xs),
            "ys": {k: list(v) for k, v in self.ys.items()},
        }

    def __repr__(self):
        return f"{self.__class__.__name__}({self.domain}, {len(self.xs)}, audiences={list(self.ys.keys())})"


def max_with_constant(threshold, curve: Curve, key: str):
    xs = [curve.xs[0]]
    ys = [max(curve.ys[key][0], threshold)]

    above_threshold = curve.ys[key][0] > threshold

    for i in range(1, len(curve)):
        y = curve.ys[key][i]

        if y > threshold and above_threshold:
            xs.append(curve.xs[i])
            ys.append(y)

        elif y > threshold and not above_threshold:
            above_threshold = True
            xs.append(
                _intersect(
                    threshold,
                    curve.xs[i - 1],
                    curve.xs[i],
                    curve.ys[key][i - 1],
                    y,
                )
            )
            xs.append(curve.xs[i])
            ys.append(threshold)
            ys.append(y)

        elif y <= threshold and above_threshold:
            above_threshold = False
            xs.append(
                _intersect(
                    threshold,
                    curve.xs[i - 1],
                    curve.xs[i],
                    curve.ys[key][i - 1],
                    y,
                )
            )
            ys.append(threshold)

    if not above_threshold:
        xs.append(curve.xs[-1])
        ys.append(threshold)

    return Curve(xs, {key: np.array(ys)})


def _intersect(threshold, x_min, x_max, y_min, y_max):
    p = threshold - y_min / (y_max - y_min)
    return x_min + p * (x_max - x_min)


def distance(curve1: Curve, curve2: Curve):
    """The maximum pointwise distance between curves."""
    xs = np.unique(np.concatenate((curve1.xs, curve2.xs)))
    y_ = np.abs(curve1(xs[0]) - curve2(xs[0]))  # type: ignore
    x_ = xs[0]
    for x in xs[1:]:
        y = np.abs(curve1(x) - curve2(x))  # type: ignore
        if y > y_:
            x_ = x
            y_ = y
    return (x_, y_)
