import itertools

import util
from solvers import solvers, solve_explicit, solve_implicit, solve_implicit_improved, solve_crank_nicolson, \
  solve_crank_nicolson_improved
from solvers.theoretical import solve_theoretical
from util import Grid

import numpy as np


def test_order(errors, next_params):
  errors = dict(errors)
  for params in list(errors.keys()):
    if params not in errors:
      continue

    # chain start
    param_chain = []
    error_chain = []

    while params in errors:
      param_chain.append(params)
      error_chain.append(errors[params])

      del errors[params]

      params = next_params(params)

    if len(param_chain) < 2:
      continue

    print(' ', ' -> '.join(map(str, param_chain)))
    e = np.array(error_chain, dtype=float)
    ed = e[:-1] / e[1:]
    print(' ', ' '.join([f'{i:6.2f}' for i in ed]))


def test_solver(solver, Is, Ks):
  param_pairs = itertools.product(Is, Ks)
  param_pairs = sorted(param_pairs, key=np.prod)

  errors_mse = {}
  errors_mae = {}
  errors_mae_last = {}

  for I, K in param_pairs:
    grid = Grid(I, K)

    solution = solver(grid)
    if solution is None:
      continue

    reference = solve_theoretical(grid)

    errors_mse[(I, K)] = util.mse(reference, solution)
    errors_mae[(I, K)] = util.mae(reference, solution)
    errors_mae_last[(I, K)] = util.mae_last(reference, solution)

  print()
  print()
  print(solver.__name__)

  print('O(h_x, h_t) MAEL       SHOULD BE >=2')
  test_order(errors_mae_last, next_params=lambda params: (params[0] * 2, params[1] * 2))

  print(f'O(h_x, h_t^2) MAEL    SHOULD BE >=4')
  test_order(errors_mae_last, next_params=lambda params: (params[0] * 4, params[1] * 2))

  print(f'O(h_x^2, h_t) MAEL    SHOULD BE >=4')
  test_order(errors_mae_last, next_params=lambda params: (params[0] * 2, params[1] * 4))


if __name__ == '__main__':
  # Is = [4, 8, 16, 32, 64, 128, 256, 512]
  # Ks = [4, 8, 16, 32, 64, 128, 256, 512]
  # Is = [4, 8, 16, 32, 64, 128, 256, 512]
  # Ks = [32, 64, 128, 256, 512, 1024, 2048, 4096]
  Is = [4, 8, 16, 32, 64]
  Ks = [32, 64, 128, 256, 512, 1024]

  test_solver(solve_explicit, Is, Ks)
  test_solver(solve_implicit, Is, Ks)
  test_solver(solve_implicit_improved, Is, Ks)
  test_solver(solve_crank_nicolson, Is, Ks)
  test_solver(solve_crank_nicolson_improved, Is, Ks)