#!/usr/bin/env python3
from qpms import TMatrixGenerator, BaseSpec, eV, hbar
import numpy as np
import sys

errors = 0

def tmg_diagonal_fun(tmatrix, omega):
    '''
    Example of a python function used as a custom T-matrix generator

    It receives a CTMatrix argument with pre-filled BaseSpec
    (in tmatrix.spec) and angular frequency.

    It has to fill in the T-matrix elements tmatrix[...] 
    (a numpy array of shape (len(tmatrix.spec),len(tmatrix.spec)))
    and return zero (on success) or other integral value on error.

    Note that this in justa an example of using the API,
    not supposed to be anything physical.
    '''
    l = tmatrix.spec.l()
    tmatrix[...] = np.diag(1./l**2)
    return 0

# Wrap the function as an actual TMatrixGenerator
tmg_diagonal = TMatrixGenerator(tmg_diagonal_fun)

bspec = BaseSpec(lMax=2)

tmatrix = tmg_diagonal(bspec, (2.0+.01j) * eV/hbar)

errors += np.sum(tmatrix[...] != np.diag(1./bspec.l()**2))

sys.exit(errors)

