'''schrodinger_solve

Solve one dimensional, time independent Schrödinger's equation for an user
provided system description '''

import sys
import argparse
from pathlib import Path

import numpy as np
from numpy.typing import NDArray

from schroedinger import (
    Config, build_potential, solve_schroedinger
)

DESCRIPTION='Solve time independent Schrödinger\'s equation for a given system.'

def save_wavefuncs(
        filename: Path,
        pos: NDArray[np.float64],
        wavefuncs: NDArray[np.float64]
) -> None:
    '''Save wave functions to file

    :param filename: Filename where to output wave functions
    :param pos: x Axis
    :param wavefuncs: Wave functions
    '''
    content = np.zeros((pos.shape[0], wavefuncs.shape[1] + 1))
    content[:, 0] = pos
    for i in range(wavefuncs.shape[1]):
        content[:, 1 + i] = wavefuncs[:, i]
    np.savetxt(filename, content)


def save_expvalues(
        filename: Path,
        pos: NDArray[np.float64],
        wavefuncs: NDArray[np.float64]
) -> None:
    '''Save expected values to file

    :param filename: Filename where to output expected values
    :param pos: x Axis
    :param wavefuncs: Wave functions
    '''
    nfuncs = wavefuncs.shape[1]
    delta = np.abs(pos[1] - pos[0])
    expvalues = np.zeros((nfuncs, 2))
    for i in range(nfuncs):
        exp_x = delta * np.sum(wavefuncs[:, i] * pos * wavefuncs[:, i])
        exp_xsq = delta * np.sum(wavefuncs[:, i] * pos**2 * wavefuncs[:, i])
        expvalues[i, 0] = exp_x
        expvalues[i, 1] = np.sqrt(exp_xsq - exp_x ** 2)
    np.savetxt(filename, expvalues)


def main() -> None:
    '''Program entry point
    '''
    parser = argparse.ArgumentParser(
        prog='schrodinger_solve',
        description=DESCRIPTION,
        epilog='')

    parser.add_argument('filename',
                        help='File describing the system to solve')
    parser.add_argument('-o', '--output-dir',
                        help='Output directory for the results')

    args = parser.parse_args()

    conf = Config(args.filename)
    output_path = Path(args.output_dir) if args.output_dir else Path.cwd()

    potential, delta = build_potential(conf)

    eig, vec = solve_schroedinger(conf.mass, potential[:, 1], delta,
                                  conf.eig_interval)

    try:
        np.savetxt(output_path / 'potential.dat', potential)
        np.savetxt(output_path / 'energies.dat', eig)
        save_wavefuncs(output_path / 'wavefuncs.dat', potential[:, 0], vec)
        save_expvalues(output_path / 'expvalues.dat', potential[:, 0], vec)
    except FileNotFoundError:
        print('Output files could not be saved.'
              ' Are you sure the output directory exists?')
        sys.exit(1)


if __name__ == '__main__':
    main()