# Module to fit a function with 3 linear parameters on given points.
# Last edited on 2021-03-07 21:47:46 by jstolfi

import rn
import rmxn
import sys
import math

def mat(G, wt):
  # Expects an {nd} by {ng} /basis matrix/ {G} where element {G[k][i]} is
  # the basis function {i} evaluated at sample argument {k}, and an
  # {nd}-element /weight table/ {wt} where {wt[k]} is the weight of
  # sample {k}.
  #
  # Returns the {ng} by {ng} /moment matrix/ {M = G'*G} of the quadratic LSQ formula
  # where {G'} is the transpose of {G}.
  # 
  nd = len(G)
  assert len(wt) == nd
  ng = len(G[0])
  M = [None]*ng
  for i in range(ng):
    Mi = [None]*ng
    for j in range(ng):
      Mij = 0
      for k in range(nd):
        Mij += G[k][i]*G[k][j]*wt[k]
      Mi[j] = Mij
    M[i] = Mi
  return M
  # ----------------------------------------------------------------------
  
def rhs(D, G, wt):
  # Assumes {g} and {wt} are as in {mat}, and that 
  # {D} is an {nd} by {nq} matrix of observations,
  # Returns the {ng} by {nq} right-hand-side matrix of the quadratic LSQ formula.
  
  nd = len(D)
  nq = len(D[0])
  assert len(G) == nd
  ng = len(G[0])
  assert len(wt) == nd
  
  B = [None]*ng
  for i in range(ng):
    Bi = [None]*nq
    for j in range(nq):
      Bij = 0
      for k in range(nd):
        Bij += D[k][j]*G[k][i]*wt[k]
      Bi[j] = Bij
    B[i] = Bi
  return B
  # ----------------------------------------------------------------------
  
def solve(M, B):
  # Solves the system {M * A = B} for an {ng} by {ng} moment matrix {M}
  # and an {ng} by {nq}  right-hand-side {B}, returns the {ng} by {nq} solution
  # matrix {A}.
  N = rmxn.inv(M)
  A = rmxn.mul(N,B)
  return A
  # ----------------------------------------------------------------------