import functools
import pickle

import numpy as np

from util import Grid

with open('./roots128.pickle', 'rb') as file:
  roots128 = pickle.load(file)


def u_elem128(x, t, alpha_n):
  x = np.float128(x)
  t = np.float128(t)
  alpha_n = np.float128(alpha_n)
  return -4 * np.pi ** 2 * (alpha_n * np.cos(alpha_n * x / 10) + 40 * np.sin(alpha_n * x / 10)) * (
      alpha_n * np.sin(2 * alpha_n) - 40 * np.cos(2 * alpha_n) + 40) * np.exp(-3 * alpha_n ** 2 * t / 500) / (
             5 * (alpha_n ** 2 - np.pi ** 2) * (alpha_n ** 2 * (4 * alpha_n + np.sin(4 * alpha_n)) - 80 * alpha_n * (
             np.cos(4 * alpha_n) - 1) + 6400 * alpha_n - 1600 * np.sin(4 * alpha_n)))


@functools.lru_cache(maxsize=1000000)
def u128(x, t, root_n=len(roots128)):
  return sum(u_elem128(x, t, roots128[:root_n]))


#@cachier(stale_after=datetime.timedelta(hours=1))
def solve_theoretical(grid: Grid) -> np.ndarray:
  ug = np.zeros((grid.K + 1, grid.I + 1), dtype=float)

  for k, t in enumerate(grid.ts):
    for i, x in enumerate(grid.xs):
      ug[k, i] = u128(x, t, 1000)

  return ug
