#! /usr/bin/python3
# Last edited on 2021-07-31 16:04:08 by jstolfi

import sys
from math import sqrt, hypot, sin, cos, exp, log, floor, ceil, inf, nan, pi;

import rn

from sampling_2d_grid_choose import sampling_2d_grid_choose
from weights_mesa import weights_mesa

from basis_2d_grid_choose import basis_2d_grid_choose
from basis_2d_hartley_choose import basis_2d_hartley_choose

from basis_2d_gauss_sample import basis_2d_gauss_sample
from basis_2d_hann_sample import basis_2d_hann_sample
from basis_2d_hartley_sample import basis_2d_hartley_sample
from basis_sampled_orthize import basis_sampled_orthize
from basis_sampled_residuals_compute import basis_sampled_residuals_compute
from basis_sampled_residuals_combine import basis_sampled_residuals_combine

from basis_sampled_write import basis_sampled_write

def main():
  # The sampling points and weights for approximation:
  ns1 = 30;     # Num samples in each axis.
  s = sampling_2d_grid_choose(-1,+1,ns1, -1,+1,ns1)
  grid = True # For {basis_sampled_write}.
  
  # The approximation basis (grid of humps):
  nb1 = int(sys.argv[1]);
  nb = nb1*nb1
  span = True
  c, r = basis_2d_grid_choose(-1,+1,nb1, -1,+1,nb1, span);
  assert len(r) == nb and len(c) == nb
  sys.stderr.write("appr basis = %d x %d = %d elements\n" % (nb1, nb1, nb))

  mrg = 1.0/(nb1+1) # Margin width for weight window.
  w_appr = weights_mesa(s, mrg)
  # w_appr = None # Approximate with Euclidean (unweighted) metric.

  # w_eval = None # Evaluate error using Euclidean metric.
  w_eval = w_appr # Evaluate error using the approximation metric.
  
  # The eval basis (Harmonics):
  fmax = 2; # Max frequency of waves.
  iso = False
  fr = basis_2d_hartley_choose(fmax,fmax,iso);
  nf = len(fr)
  sys.stderr.write("eval basis = %d elements (fmax = %d iso = %s)\n" % (nf, fmax, iso))
  
  # Sample the two bases:
  # B = basis_2d_hann_sample(c, r, s);
  B = basis_2d_gauss_sample(c, r, s);
  wr = open("out/B.txt", "w")
  basis_sampled_write(wr, s, B, w_appr, grid)
  
  # Orthogonalize {B} with the projection norm (weights {w}) yielding {C}:
  C = basis_sampled_orthize(B, w_appr)
  wr = open("out/C.txt", "w")
  basis_sampled_write(wr, s, C, w_appr, grid)

  F = basis_2d_hartley_sample(fr, s);
  wr = open("out/F.txt", "w")
  basis_sampled_write(wr, s, F, w_appr, grid)
  
  # Project each {F} basis element on {<C>} with the projection yielding
  # the coefficient matrix {M} and the residual vector {R}:
  R = basis_sampled_residuals_compute(F, C, w_appr)
  wr = open("out/R.txt", "w")
  basis_sampled_write(wr, s, R, w_appr, grid)
  
  # Combine all residuals into an error maps:
  e = basis_sampled_residuals_combine(R, None)
  
  # Plot the error map:
  wr = open("out/emap.txt", "w")
  basis_sampled_write(wr, s, [ e ], w_appr, grid)
  return 0
  # ----------------------------------------------------------------------

main()