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)


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

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

   @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 = NonLinear.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 at [{bounds[0]},{bounds[1]}]", "or")
      plt_append(
          sp, *(sol2[1]), f"bisect at [{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 = NonLinear.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 at [{bounds[0]},{bounds[1]}]", "or")
      plt_append(
          sp, *(sol2[1]), f"newton at [{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 = "np.sin(x) * np.sqrt(np.abs(x))"

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


if __name__ == "__main__":
   main()