import numpy as np
import scipy.sparse
import scipy.sparse.linalg

import consts
from util import Grid


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

  ug[0, :] = np.sin(np.pi * grid.xs / consts.l) ** 2 * 4 / 10

  for k in range(1, grid.K + 1):
    diags = {
      -1: np.zeros(grid.I + 1, dtype=float),
      0: np.zeros(grid.I + 1, dtype=float),
      +1: np.zeros(grid.I + 1, dtype=float),
    }
    b = np.zeros(grid.I + 1, dtype=float)

    # left
    diags[0][0] = consts.H * grid.h_x + 1
    diags[+1][0] = -1

    # middle
    for i in range(1, grid.I):
      diags[-1][i] = grid.gamma
      diags[0][i] = - 2 * grid.gamma - 2
      diags[+1][i] = grid.gamma
      b[i] = -(
          ug[k - 1][i - 1] * grid.gamma +
          ug[k - 1][i] * (-2 * grid.gamma + 2) +
          ug[k - 1][i + 1] * grid.gamma
      )

    # right
    diags[0][grid.I] = consts.H * grid.h_x + 1
    diags[-1][grid.I] = -1

    A = scipy.sparse.diags([
      diags[-1][1:],
      diags[0],
      diags[+1][:-1],
    ], [-1, 0, 1], format='csc')

    ug[k, :] = scipy.sparse.linalg.spsolve(A, b)

  return ug


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

  ug[0, :] = np.sin(np.pi * grid.xs / consts.l) ** 2 * 4 / 10

  for k in range(1, grid.K + 1):
    diags = {
      -1: np.zeros(grid.I + 1, dtype=float),
      0: np.zeros(grid.I + 1, dtype=float),
      +1: np.zeros(grid.I + 1, dtype=float),
    }
    b = np.zeros(grid.I + 1, dtype=float)

    # left
    diags[0][0] = consts.H * grid.gamma * grid.h_x + grid.gamma + 1
    diags[+1][0] = - grid.gamma
    b[0] = (
        (- consts.H * grid.gamma * grid.h_x - grid.gamma + 1) * ug[k - 1, 0] +
        grid.gamma * ug[k - 1, 1]
    )

    # middle
    for i in range(1, grid.I):
      diags[-1][i] = grid.gamma
      diags[0][i] = - 2 * grid.gamma - 2
      diags[+1][i] = grid.gamma
      b[i] = -(
          ug[k - 1][i - 1] * grid.gamma +
          ug[k - 1][i] * (-2 * grid.gamma + 2) +
          ug[k - 1][i + 1] * grid.gamma
      )

    # right
    diags[0][grid.I] = consts.H * grid.gamma * grid.h_x + grid.gamma + 1
    diags[-1][grid.I] = - grid.gamma
    b[grid.I] = (
        (- consts.H * grid.gamma * grid.h_x - grid.gamma + 1) * ug[k - 1, grid.I] +
        grid.gamma * ug[k - 1, grid.I - 1]
    )

    A = scipy.sparse.diags([
      diags[-1][1:],
      diags[0],
      diags[+1][:-1],
    ], [-1, 0, 1], format='csc')

    ug[k, :] = scipy.sparse.linalg.spsolve(A, b)

  return ug
