from dataclasses import dataclass, field
import numpy as np

import consts


@dataclass
class Grid:
  I: int
  K: int

  xs: np.ndarray = field(repr=False, compare=False)
  ts: np.ndarray = field(repr=False, compare=False)

  h_x: float = field(repr=False, compare=False)
  h_t: float = field(repr=False, compare=False)

  gamma: float = field(repr=False, compare=False)

  def __init__(self, I: int, K: int):
    self.I = I
    self.K = K

    self.xs = np.linspace(0, consts.l, I + 1)
    self.ts = np.linspace(0, consts.T, K + 1)

    self.h_x = consts.l / I
    self.h_t = consts.T / K

    self.gamma = consts.D * self.h_t / (self.h_x ** 2)

  def __hash__(self):
    return hash((self.I, self.K))


def mse(ug1: np.ndarray, ug2: np.ndarray) -> float:
  return np.mean((ug1 - ug2) ** 2)


def mae(ug1: np.ndarray, ug2: np.ndarray) -> float:
  return np.max(np.abs(ug1 - ug2))


def mae_last(ug1: np.ndarray, ug2: np.ndarray) -> float:
  return np.max(np.abs(ug1[-1, :] - ug2[-1, :]))
