import sys
import argparse
from pathlib import Path

import numpy as np

from schroedinger import (
    Config, potential_interp, build_potential, solve_schroedinger

)

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

def save_wavefuncs(filename, x, v):
    wavefuncs = np.zeros((x.shape[0], v.shape[1] + 1))
    wavefuncs[:, 0] = x
    for i in range(v.shape[1]):
        wavefuncs[:, 1 + i] = v[:, i]
    np.savetxt(filename, wavefuncs)


def save_expvalues(filename, x, v):
    n = v.shape[1]
    delta = np.abs(x[1] - x[0])
    expvalues = np.zeros((n, 2))
    for i in range(n):
        exp_x = delta * np.sum(v[:, i] * x * v[:, i])
        exp_xsq = delta * np.sum(v[:, i] * x**2 * v[:, i])
        expvalues[i, 0] = exp_x
        expvalues[i, 1] = np.sqrt(exp_xsq - exp_x ** 2)
    np.savetxt(filename, expvalues)


def main():
    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)

    e, v = solve_schroedinger(conf.mass, potential[:, 1], delta,
                              conf.eig_interval)

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


if __name__ == '__main__':
    main()