'''schroedinger.py

Collection of method common to solve and tests
'''

import sys
from pathlib import Path
from typing import TextIO, Type

import numpy as np
from numpy.polynomial import Polynomial
from numpy.typing import NDArray

from scipy.interpolate import CubicSpline
from scipy.linalg import eigh_tridiagonal


Interpolator = Type[callable] | Type[Polynomial] | Type[CubicSpline]
# type Interpolator = callable | Polynomial | CubicSpline


class Config:
    '''Wrapper for reading, parsing and storing the contents of a configuration
    file
    '''
    def __init__(self, path):
        current_line = 0

        # Ensure Path object
        if path is not Path:
            path = Path(path)

        def next_parameter(file: TextIO):
            '''Read next parameter, ignoring comments or empty lines'''
            content = None
            nonlocal current_line
            while not content:
                line = file.readline()
                current_line += 1
                index = line.find('#')
                content = line[0:index].strip()
            return content

        with open(path, 'r', encoding='utf8') as file:
            try:
                self.mass = float(next_parameter(file))
                start, end, steps = next_parameter(file).split()
                self.interval = [float(start), float(end), int(steps)]
                self.eig_interval = [int(attr) - 1 for attr in next_parameter(file).split()]
                self.interpolation = next_parameter(file)

                npoints = int(next_parameter(file))
                self.points = np.zeros((npoints, 2))
                for i in range(npoints):
                    line = next_parameter(file)
                    self.points[i] = np.array([float(comp) for comp in line.split()])
            except ValueError:
                print(f'Syntax error in \'{path.name}\' line {current_line}')
                sys.exit(1)


def potential_interp(
        interpolation: str,
        points: NDArray[np.float64]
) -> Interpolator:
    '''Create an interpolator for a set of predefined potential points

    :param interpolation: Kind of interpolation
    :param points: Points to interpolate within
    :return: Interpolating object
    '''
    interpolator = None

    if interpolation == 'linear':
        def line(pos):
            return np.interp(pos, points[:, 0], points[:, 1])
        interpolator = line
    elif interpolation == 'polynomial':
        interpolator = Polynomial.fit(points[:, 0], points[:, 1],
                                      points.shape[0] - 1)
    elif interpolation == 'cspline':
        interpolator = CubicSpline(points[:, 0], points[:, 1])
    else:
        raise ValueError('Invalid interpolator kind. Use any of: linear, polynomial or cspline')

    return interpolator


def build_potential(
        config: Config
) -> tuple[NDArray[np.float64], np.float64]:
    '''Build a potential based on the options inside a configuration file

    :param config: System parameters
    :return: Potential and distance between samples
    '''
    start, end, steps = config.interval
    potential = np.zeros((steps, 2))
    potential[:, 0] = np.linspace(start, end, steps)
    delta = np.abs(potential[1, 0] - potential[0, 0])
    interp = potential_interp(config.interpolation, config.points)
    potential[:, 1] = interp(potential[:, 0])
    return potential, delta


def solve_schroedinger(
        mass: float,
        potential: NDArray[np.float64],
        delta: float,
        eig_interval: tuple[int, int]=None
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    '''Solve the one dimensional, time independent Schrödinger's equation for a
    particle of given mass inside a discretized potential

    :param mass: Mass of the particle
    :param potential: Discretized potential
    :param delta: Distance between points in the potential
    :param eig_interval: Interval of quantum numbers for which the states are to be calculated
    :return: Eigenvalues and normalized wave functions
    '''
    qnumber = potential.shape[0]
    coeff = 1 / mass / delta**2
    eig, vec = eigh_tridiagonal(coeff + potential,
                                -coeff * np.ones(qnumber - 1, dtype=np.float64) / 2.0,
                                select='i', select_range=eig_interval)
    # Normalize eigenfunctions
    for i in range(eig.shape[0]):
        vec[:, i] /= np.sqrt(delta * np.sum(np.abs(vec[:, i])**2))

    return eig, vec