import scipy.integrate as sitg
import scipy.interpolate as sitp
import scipy.optimize as sopt
import scipy.linalg as salg
import math as m
import numpy as np

import matplotlib.pyplot as plt


def create_subplot():
   return plt.subplots(layout='constrained')[1]


def plt_append(sp, x: list[float], y: list[float], label: str, format: str):
   sp.plot(x, y, format, label=label)


def generate_array(min, max, density=10):
   point_count = int(m.fabs(max-min)*density)
   x = np.linspace(min, max, point_count)
   return list(x.tolist())


class NonLinear:
   bisect_exp = "x**2 * np.sin(x)"
   newton_exp = "np.sin(x) * np.sqrt(np.abs(x))"

   @staticmethod
   def slice_array(range: list[float], val_min, val_max):
      def index_search(range: list[float], val):
         i = 0
         for v in range:
            if v >= val:
               return i
            i += 1
         return -1

      index_l = index_search(
          range, val_min) if val_min is not None else range.index(min(range))
      index_r = index_search(
          range, val_max) if val_max is not None else range.index(max(range))
      return range[index_l:index_r+1]

   @staticmethod
   def bisect(x, x_min, x_max):
      def f(x): return eval(NonLinear.bisect_exp)
      y = f(np.array(x))
      root = sopt.bisect(f, x_min, x_max)
      solution = root[0] if root is tuple else root
      return list(y), (float(solution), float(f(solution)))

   @staticmethod
   def plot_bisect():
      bounds = 0, 6
      split_val = 1
      x1 = generate_array(bounds[0], bounds[1])
      x2 = NonLinear.slice_array(x1, split_val, None)

      sp = create_subplot()

      sol1 = NonLinear.bisect(x1, bounds[0], bounds[1])
      sol2 = NonLinear.bisect(x2, split_val, bounds[1])

      plt_append(
          sp, x1, sol1[0], f"Исходные данные (y={NonLinear.bisect_exp})", "-b")
      plt_append(
          sp, *(sol1[1]), f"bisect на [{bounds[0]},{bounds[1]}]", "or")
      plt_append(
          sp, *(sol2[1]), f"bisect на [{split_val},{bounds[1]}]", "og")

      sp.set_title("scipy.optimize.bisect")
      sp.legend(loc='lower left')

   @staticmethod
   def newton(x, x0):
      def f(x): return eval(NonLinear.bisect_exp)
      y = f(np.array(x))
      root = sopt.newton(f, x0)
      solution = root[0] if root is tuple else root
      return list(y), (float(solution), float(f(solution)))

   @staticmethod
   def plot_newton():
      bounds = -2, 7
      split_l, split_r = 2, 5
      x1 = generate_array(bounds[0], bounds[1])
      x2 = NonLinear.slice_array(x1, split_l, split_r)
      x0_1, x0_2 = 1/100, 4
      sp = create_subplot()

      sol1 = NonLinear.newton(x1, x0_1)
      sol2 = NonLinear.newton(x2, x0_2)

      plt_append(
          sp, x1, sol1[0], f"Исходные данные (y={NonLinear.newton_exp})", "-b")
      plt_append(
          sp, *(sol1[1]), f"newton на отрезке [{bounds[0]},{bounds[1]}]", "or")
      plt_append(
          sp, *(sol2[1]), f"newton на отрезке [{split_l},{bounds[1]}]", "og")

      sp.set_title("scipy.optimize.newton")
      sp.legend(loc='lower left')

   @staticmethod
   def plot(method: str = "all"):
      if method in ["bisect", "all"]:
         NonLinear.plot_bisect()
      if method in ["newton", "all"]:
         NonLinear.plot_newton()
      plt.ylabel("y")
      plt.xlabel("x")
      plt.show()


class SLE:
   gauss_data = ([[13, 2], [3, 4]], [1, 2])
   invmatrix_data = ([[13, 2], [3, 4]], [1, 2])
   tridiagonal_data = ([[4,  5,  6,  7, 8, 9],
                       [2,  2,  2,  2, 2, 0]],
                       [1, 2, 2, 3, 3, 3])

   @staticmethod
   def var_str(index):
      return f"x{index+1}"

   @staticmethod
   def print_solution(data: list[float]):
      print("  ", end='')
      for i, val in enumerate(data[:-1]):
         print(f"{SLE.var_str(i)} = {round(val,3)}, ", end='')
      print(f"{SLE.var_str(len(data)-1)} = {round(data[-1],3)}")

   @staticmethod
   def print_data(data: tuple[list[list[float]], list[float]], tridiagonal: bool = False):
      if tridiagonal:
         new_data = []
         new_len = len(data[0][0])
         zipped = list(zip(*tuple(data[0])))
         zipped[len(zipped)-1] = (zipped[len(zipped)-1]
                                  [0], zipped[len(zipped)-2][1])
         complement_to = new_len - len(zipped[0])
         for i, val in enumerate(zipped):
            zero_r = complement_to - i
            if zero_r <= 0:
               zero_r = 0
            mid_val = list(reversed(val[1:])) + list(val)
            mid_end = len(mid_val) if zero_r > 0 else len(
                mid_val) + (complement_to - i)
            mid_beg = len(mid_val) - (new_len - zero_r) if zero_r > 0 else 0
            mid_beg = mid_beg if mid_beg >= 0 else 0
            zero_l = new_len - (zero_r + (mid_end - mid_beg))
            tmp = [0] * zero_l + \
                mid_val[mid_beg:mid_end] + [0] * zero_r
            new_data.append(tmp)
         data = (new_data, data[1])
      for i, val in enumerate(data[0]):
         print("  ", end='')
         for i_coef, coef in enumerate(val[:-1]):
            if coef != 0:
               print(f"({coef}{SLE.var_str(i_coef)}) + ", end='')
            else:
               print(f"  {coef}   + ", end='')
         print(f"({val[-1]}{SLE.var_str(len(val)-1)})", end='')
         print(f" = {data[1][i]}")

   @staticmethod
   def gauss(system: list[list[float]], b: list[float]):
      lup = salg.lu_factor(system)
      solution = salg.lu_solve(lup, b)
      return solution

   @staticmethod
   def invmatrix(system: list[list[float]], b: list[float]):
      m_inv = salg.inv(system)
      solution = m_inv @ b
      return solution

   @staticmethod
   def tridiagonal(system: list[list[float]], b: list[float]):
      solution = salg.solveh_banded(system, b, lower=True)
      return solution

   @staticmethod
   def print_gauss():
      print("Gauss method (LU decomposition)")
      print("  Input system:")
      SLE.print_data(SLE.gauss_data)
      print("  Solution:")
      SLE.print_solution(SLE.gauss(*SLE.gauss_data))

   @staticmethod
   def print_invmatrix():
      print("Inverted matrix method")
      print("  Input system:")
      SLE.print_data(SLE.invmatrix_data)
      print("  Solution:")
      SLE.print_solution(SLE.invmatrix(*SLE.invmatrix_data))

   @staticmethod
   def print_tridiagonal():
      print("Tridiagonal matrix method (Thomas algorithm)")
      print("  Input system:")
      SLE.print_data(SLE.tridiagonal_data, True)
      print("  Solution:")
      SLE.print_solution(SLE.tridiagonal(*SLE.tridiagonal_data))

   @staticmethod
   def print(method="all"):
      if method in ["gauss", "all"]:
         SLE.print_gauss()
      if method in ["invmatrix", "all"]:
         SLE.print_invmatrix()
      if method in ["banded", "all"]:
         SLE.print_tridiagonal()


class Approx:
   function_exp = "np.sin(x) * np.sqrt(np.abs(x))"
   least_sq_exp = "np.sin(x) * np.abs(x)"

   @staticmethod
   def get_function_exp_der(*args):
      function_der_exp = "(x * np.sin(x) + 2 * x**2 * np.cos(x)) / (2 * np.sqrt(np.abs(x)) ** 3)"
      result = ()
      for i in args:
         array = []
         for x in i:
            array.append(eval(function_der_exp))

         result = result + (array,)
      return result

   @staticmethod
   def generate_y(x_array, function):
      result = []
      for x in x_array:
         result.append(eval(function))
      return result

   @staticmethod
   def lagrange(x, y):
      return sitp.lagrange(x, y)

   @staticmethod
   def get_approx_data(function=function_exp, bounds=[-6, 6]):
      x1 = generate_array(bounds[0], bounds[1], 1/2)
      x2 = generate_array(bounds[0], bounds[1], 1)
      y1 = Approx.generate_y(x1, function)
      y2 = Approx.generate_y(x2, function)
      x_real = generate_array(bounds[0], bounds[1])
      y_real = Approx.generate_y(x_real, function)
      return x1, x2, y1, y2, x_real, y_real

   @staticmethod
   def plot_lagrange():
      x1, x2, y1, y2, x_real, y_real = Approx.get_approx_data()

      sp = create_subplot()
      sol1 = np.polynomial.polynomial.Polynomial(
          Approx.lagrange(x1, y1).coef[::-1])
      sol2 = np.polynomial.polynomial.Polynomial(
          Approx.lagrange(x2, y2).coef[::-1])

      plt_append(
          sp, x_real, y_real, f"Исходные данные (y={Approx.function_exp})", "--b")
      plt_append(
          sp, x_real, sol1(np.array(x_real)), f"f1 = lagrange, кол-во точек = {len(x1)}", "-m")
      plt_append(
          sp, x_real, sol2(np.array(x_real)), f"f2 = lagrange, кол-во точек = {len(x2)}", "-r")
      plt_append(
          sp, x1, y1, f"Исходные точки для f1", ".m")
      plt_append(
          sp, x2, y2, f"Исходные точки для f2", ".r")

      sp.set_title("scipy.interpolate.lagrange")
      sp.legend(loc='lower left')

   @staticmethod
   def plot_spline():
      x1, x2, y1, y2, x_real, y_real = Approx.get_approx_data()
      d1, d2 = Approx.get_function_exp_der(x1, x2)

      for interpolator in [sitp.CubicSpline,
                           sitp.PchipInterpolator,
                           sitp.CubicHermiteSpline,
                           sitp.Akima1DInterpolator]:
         sp = create_subplot()

         if interpolator.__name__ != "CubicHermiteSpline":
            args1 = x1, y1
            args2 = x2, y2
         else:
            args1 = x1, y1, d1
            args2 = x2, y2, d2

         sol1 = interpolator(*args1)
         sol2 = interpolator(*args2)
         plt_append(
             sp, x_real, y_real, f"Исходные данные (y={Approx.function_exp})", "--b")
         plt_append(
             sp, x_real, sol1(np.array(x_real)), f"f1 = {interpolator.__name__}, кол-во точек = {len(x1)}", "-m")
         plt_append(
             sp, x_real, sol2(np.array(x_real)), f"f2 = {interpolator.__name__}, кол-во точек = {len(x2)}", "-r")
         plt_append(
             sp, x1, y1, f"Исходные точки для f1", ".m")
         plt_append(
             sp, x2, y2, f"Исходные точки для f2", ".r")

         sp.set_title(f"scipy.interpolate.{interpolator.__name__}")
         sp.legend(loc='lower left')

   @staticmethod
   def linear(x, a, b):
      return a*x + b

   @staticmethod
   def quadratic(x, a, b, c):
      return a * (x**2) + (b*x) + c

   @staticmethod
   def fract(x, a, b, c):
      return x / (a * x + b) - c

   @staticmethod
   def noise_y(y, rng):
      diff = max(y) - min(y)
      noise_coeff = diff*(10/100)
      return y + (noise_coeff * rng.normal(size=len(y)))

   @staticmethod
   def plot_least_squares_curvefit():
      rng = np.random.default_rng()
      bounds = [3, 6]
      x1, x2, y1, y2, x_real, y_real = Approx.get_approx_data(
          Approx.least_sq_exp, bounds)
      x_real = np.array(x_real)

      y_real = Approx.noise_y(y_real, rng)
      base_functions = [Approx.linear,
                        Approx.quadratic, (Approx.fract, "x/(ax+b)")]

      sp = create_subplot()
      plt_append(
          sp, x_real, y_real, f"y={Approx.least_sq_exp} на [{bounds[0]};{bounds[1]}], с шумом", ".b")
      for bf in base_functions:
         if isinstance(bf, tuple):
            bf, desc = bf[0], bf[1]
         else:
            bf, desc = bf, None
         optimal_params, _ = sopt.curve_fit(bf, x_real, y_real)
         desc_str = f" ({desc}) " if desc is not None else ""
         plt_append(
             sp, x_real, bf(np.array(x_real), *optimal_params),
             f"МНК, вид функции - {bf.__name__}{desc_str}", "-")
      sp.set_title(f"scipy.optimize.curve_fit")
      sp.legend(loc='lower left')

   @staticmethod
   def plot_least_squares():
      rng = np.random.default_rng()

      def exponential(x, a, b, c):
         return np.sin(x) * np.sqrt(np.abs(x))
      exponential.str = "np.sin(x) * np.sqrt(np.abs(x))"

      def gen_y(x, a, b, c, noise=0., n_outliers=0):
         y = exponential(x, a, b, c)
         error = noise * rng.standard_normal(x.size)
         outliers = rng.integers(0, x.size, n_outliers)
         error[outliers] *= 10
         return y + error

      def loss(params, x, y):
         return (exponential(x, params[0], params[1], params[2])) - y
      params0 = np.array([0.1, 1, 0])
      bounds = [-5, 3]
      params_real = (3, 1, 5)
      x_approx = np.array(generate_array(bounds[0], bounds[1], 4))
      y_approx = np.array(gen_y(x_approx, *params_real,
                                noise=0.3, n_outliers=4))

      params_lsq = sopt.least_squares(
          loss, params0, loss='linear', args=(x_approx, y_approx)).x
      params_soft_l1 = sopt.least_squares(
          loss, params0, loss='soft_l1', args=(x_approx, y_approx),f_scale=0.1).x
      params_cauchy = sopt.least_squares(
          loss, params0, loss='cauchy', args=(x_approx, y_approx), f_scale=2).x

      x_real = np.array(generate_array(bounds[0], bounds[1]))
      y_real = np.array(gen_y(x_real, *params_real, 0, 0))

      sp = create_subplot()
      sp.plot(x_real, y_real, "-b",
              label=f"y={exponential.str} на [{bounds[0]};{bounds[1]}]")
      sp.plot(x_approx, y_approx, ".r", label=f"Табличные значения с шумом")
      sp.plot(x_real, gen_y(x_real, *params_lsq), color="green",
              label=f"loss=\"linear\"", linestyle=(0, (5, 10)))
      sp.plot(x_real, gen_y(x_real, *params_soft_l1), color="magenta",
              label=f"loss=\"soft_l1\"", linestyle=(5, (5, 10)))
      sp.plot(x_real, gen_y(x_real, *params_cauchy), color="black",
              label=f"loss=\"cauchy\"", linestyle=(7, (5, 10)))

      sp.set_title(f"scipy.optimize.least_squares")
      sp.legend(loc='lower left')

   @staticmethod
   def plot(method: str = "all"):
      if method in ["lagrange", "all"]:
         Approx.plot_lagrange()
      if method in ["spline", "all"]:
         Approx.plot_spline()
      if method in ["least_squares_curvefit", "all"]:
         Approx.plot_least_squares_curvefit()
      if method in ["least_squares", "all"]:
         Approx.plot_least_squares()
      plt.ylabel("y")
      plt.xlabel("x")
      plt.show()


def main():
   NonLinear.plot()
   SLE.print()
   Approx.plot()


if __name__ == "__main__":
   main()