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

import consts


#@cachier(stale_after=datetime.timedelta(hours=1))
def solve_implicit(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 - 1
      diags[+1][i] = grid.gamma
      b[i] = -ug[k - 1, i]

    # right
    diags[-1][grid.I] = -1
    diags[0][grid.I] = consts.H * grid.h_x + 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_implicit_improved(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] = (
        2 * consts.H * grid.gamma * grid.h_x +
        2 * grid.gamma +
        1
    )
    diags[+1][0] = - 2 * grid.gamma
    b[0] = ug[k - 1, 0]

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

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

    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