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]


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 plt_append(sp, x: list[float], y: list[float], label: str, format: str):
      sp.plot(x, y, format, label=label)

   @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])

      NonLinear.plt_append(
          sp, x1, sol1[0], f"Исходные данные (y={NonLinear.bisect_exp})", "-b")
      NonLinear.plt_append(
          sp, *(sol1[1]), f"bisect at [{bounds[0]},{bounds[1]}]", "or")
      NonLinear.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)

      NonLinear.plt_append(
          sp, x1, sol1[0], f"Исходные данные (y={NonLinear.newton_exp})", "-b")
      NonLinear.plt_append(
          sp, *(sol1[1]), f"newton at [{bounds[0]},{bounds[1]}]", "or")
      NonLinear.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()


def main():
   NonLinear.plot()


if __name__ == "__main__":
   main()