#!python
from __future__ import print_function
from __future__ import division

import sys, os

import numpy as np
import matplotlib.pyplot as plt
from cycler import cycler

import cellconstructor as CC
import cellconstructor.Phonons

import sscha
import tdscha, tdscha.DynamicalLanczos as DL
import sscha.Ensemble
MSG = """
TDSCHA  
======

Analyze the convergence with the number of steps of a Lanczos
calculation
Usage:

tdscha-convergence-analysis  lanczos_status_file [smearing]

lanczos_status_file is either a npz or a abc file generated by the script.
text output file can be converted into the abc file with the tdscha-output2abc.py utility.

The smearing is in cm-1
"""

MSG_OUTPUT2ABC = """
TDSCHA  
======

Convert the output file of a Lanczos calculation 
into a .abc file to analyze easily the results.

Usage:

tdscha-output2abc  outputfile  abcfile

outputfile must contain the stdout of the Lanczos calculation
abcfile is generated from the output
"""

MSG_PLOT = """
TDSCHA  
======

Plot the spectrum of a TDSCHA calculation.

Usage: 

tdscha-plot file.abc [w_start w_end [smearing]]

Pass a .abc or .npz file resulting from a linear response calculation.
Optionally you can pass a range of frequencies (cm-1) and the smearing.

"""


MSG_HESSIAN_CONV = """
TDSCHA  
======

Plot the convergence of a full hessian calculation with the Lanczos algorithm.

Usage: 

    tdscha-hessian-convergence directory prefix [initial_status] [dynamical_matrix]


Arguments: 
- directory: the directory where the lanczos steps are stored
- prefix: the prefix of the lanczos steps files
- initial_status: (optional) the status file to initialize the system. A prefix of two files, one json and one npz.
- dynamical_matrix: (optional) the dynamical matrix to use as reference for the phonon modes.

"""

def plot_hessian_convergence():
    print(MSG_HESSIAN_CONV)

    nargs = len(sys.argv)

    if nargs != 3 and nargs != 4 and nargs != 5:
        print("Wrong number of arguments.")
        exit(1)

    datadir = sys.argv[1]
    prefix = sys.argv[2]

    static_hessian = sscha.StaticHessian.StaticHessian()
    
    if nargs >= 4:
        if not os.path.exists(sys.argv[3] + ".json"):
            static_hessian.lanczos.load_status(sys.argv[3])
            static_hessian.preconitioned = False
            print("Neglecting json file")
        else:
            static_hessian.load_status(sys.argv[3])
            print("Loading from json file")

    else:
        static_hessian.lanczos.load_status(os.path.join(datadir, prefix))

    reference = None

    if nargs == 5:
        dname = os.path.dirname(sys.argv[4])
        fname = os.path.split(sys.argv[4])[-1]
        nqirr = len([x for x in os.listdir(dname) if x.startswith(fname)])
        reference = CC.Phonons.Phonons(sys.argv[4], nqirr)
    
    directory = datadir
    sh = static_hessian
        
    

    print(MSG_HESSIAN_CONV)

    # Get all the files from the minimization
    all_files = [x for x in os.listdir(directory) if x.startswith(prefix) and x.endswith(".dat") and len(x) == len(prefix) + 14]
    all_files.sort()
    finalfile = os.path.join(directory, prefix)
    if os.path.exists(finalfile):
        all_files.append(prefix)
        print("The algorithm converged!")
    else:
        print("The algorithm did not converged!")

    ws = np.zeros((sh.lanczos.pols.shape[0], len(all_files)))
    
    print()
    for i, fname in enumerate(all_files):
        sh.vector = np.loadtxt(os.path.join(directory, fname))

        G = sh.retrieve_hessian(noq = True) / np.sqrt(np.outer(sh.lanczos.m, sh.lanczos.m))

        w2 = np.linalg.eigvalsh(G)
        ws[:, i] = np.sign(w2) * np.sqrt(np.abs(w2))

        if i % 10 == 0:
            sys.stdout.write("\rProgress {:d} %".format((i * 100) // len(all_files)))
            sys.stdout.flush()
    print()
    
    # Plot
    plt.rcParams["font.family"] = "Liberation Serif"
    plt.figure(dpi = 120)
    
    for i in range(ws.shape[0]):
        plt.plot(ws[i, :] * CC.Units.RY_TO_CM)

    if reference is not None:
        w, p = reference.DiagonalizeSupercell()
        trans = CC.Methods.get_translations(p, reference.structure.generate_supercell(reference.GetSupercell()).get_masses_array())
        w = w[~trans]

        for wx in w:
            plt.axhline(wx * CC.Units.RY_TO_CM, 0, 1, ls = "--", color = "k")

    plt.xlabel("Optimization steps")
    plt.ylabel("Hessian frequencies [cm-1]")
    plt.title("Hessian convergence")
    plt.tight_layout()
    plt.show()
 

def plot():
    print(MSG_PLOT)
    if len(sys.argv) not in [2, 4, 5]:
        print("Error, wrong number of arguments.")
        exit()
    
    fname = sys.argv[1]
    assert os.path.exists(fname), "Error, file {} does not exist".format(fname)
    

    lanc = DL.Lanczos()

    print("Loading file {}".format(fname))

    if fname.endswith(".abc"):
        lanc.load_abc(fname)
    elif fname.endswith(".npz"):
        lanc.load_status(fname)
    else:
        print("ERROR, the specified file must either be a .abc, or .npz file.")
    

    w_start = 0
    w_end = 5000
    n_w = 50000
    smearing = 5

    if len(sys.argv) > 2:
        w_start = float(sys.argv[2])
        w_end = float(sys.argv[3])
    if len(sys.argv) == 5:
        smearing = float(sys.argv[4])
    
    w = np.linspace(w_start, w_end, n_w)
    w_ry = w / CC.Units.RY_TO_CM
    smearing /= CC.Units.RY_TO_CM

    gf = lanc.get_green_function_continued_fraction(w_ry, smearing = smearing, use_terminator=False)
    spectrum = - np.imag(gf)

    # Print some info about the calculation
    print()
    print("Number of poles: {}".format(len(lanc.a_coeffs)))
    
    plt.plot(w, spectrum)
    plt.xlabel("Frequency [cm-1]")
    plt.ylabel("Spectrum [a.u.]")
    plt.tight_layout()
    plt.show()


def convert():
    print(MSG_OUTPUT2ABC)
    if not len(sys.argv) == 3:
        print("Error, wrong number of arguments.")
        exit()

    output = sys.argv[1]
    abcfile = sys.argv[2]
 
    reading = False
    index = 0

    a = []
    b = []
    c = []   
    with open(output, "r") as fp:
        for line in fp.readlines():
            line = line.strip()

            if "LANCZOS ALGORITHM" in line:
                reading = True

            data = line.split()
            if len(data) != 3:
                continue

            if "a_{:d}".format(index) == data[0]:
                a.append(float(data[2]))
            if "b_{:d}".format(index) == data[0]:
                b.append(float(data[2]))
            if "c_{:d}".format(index) == data[0]:
                c.append(float(data[2]))
                index += 1

    ERROR = """
Error, the length of a, b and c does not match!
len(a) = {}; len(b) = {}, len(c) = {}
""".format(len(a), len(b), len(c))
    
    assert len(a) == len(b), ERROR 
    assert len(a) == len(c), ERROR

    # Save the abc file
    np.savetxt(abcfile, np.transpose([a, b, c]))



def tdscha_convergence_analysis():
    print(MSG)
    # Get from command line the last lanczos step
    if not len(sys.argv) in [2, 3]:
        print("Error, require the Lanczos .npz (or .abc) file to be analyzed, and (optionally) the smearing [cm-1]")
        raise ValueError("Error, I require the Lanczos .npz (or .abc) file to be analyzed)")

    fname = sys.argv[1]

    smearing = 1
    if len(sys.argv) == 3:
        smearing = float(sys.argv[2])

    # Check if the file exists
    if not os.path.exists(fname):
        raise IOError("Error, the specified file '{}' does not exist.".format(fname))

    # Check how many files are here
    #files = [x for x in  os.listdir(".") if "LANCZOS_STEP" in x and ".npz" in x]


    data = []
    nat = 0
    print ("Reading the lanczos file...")
    data = DL.Lanczos()

    if fname.endswith('.abc'):
        data.load_abc(fname)
    else:
        data.load_status(fname)
    nat = data.nat
    N_iters = len(data.a_coeffs) - 1

    # Now get the static converge
    W_START = 0
    W_END = 10000 / CC.Units.RY_TO_CM
    NW = 10000
    SMEARING = smearing / CC.Units.RY_TO_CM

    print ("Found {} iterations.".format(N_iters + 1))
    freqs = np.zeros((N_iters, 3*nat))
    dynamical = np.zeros((N_iters, NW))
    dynamical_noterm = np.zeros((N_iters, NW))
    w_array = np.linspace(W_START, W_END, NW)
    w_static = np.zeros((N_iters), dtype = np.double)

    a_coeffs = np.copy(data.a_coeffs)
    b_coeffs = np.copy(data.b_coeffs)
    c_coeffs = np.copy(data.c_coeffs)

    for i in range(N_iters):
        data.a_coeffs = a_coeffs[:i+1]
        data.b_coeffs = b_coeffs[:i]
        data.c_coeffs = c_coeffs[:i]
        
        #fc_odd = data.get_static_odd_fc(False)
        #fc_odd /= np.sqrt(np.outer(data.m, data.m))
        #w, p = np.linalg.eigh(fc_odd)
        #sign_mask = np.sign(w)
        #w = sign_mask * np.sqrt(np.abs(w))
        #w *= CC.Phonons.RY_TO_CM
        #freqs[i, :] = w 

        gf = data.get_green_function_continued_fraction(w_array, use_terminator = False, smearing = SMEARING)

        dynamical_noterm[i, :] = -np.imag( gf)
        w2 = 1 / np.real(gf[0])
        w_static[i] = np.sign(w2) * np.sqrt(np.abs(w2)) * CC.Units.RY_TO_CM
        dynamical[i, :] = -np.imag( data.get_green_function_continued_fraction(w_array, use_terminator = True, smearing = SMEARING))
        



        if i % 10 == 0:
            sys.stderr.write("\rProgress %% %d" % (i * 100 / N_iters))
            sys.stderr.flush()
    sys.stderr.write("\n")

    print ("Plotting the results...")

    plt.figure(dpi = 160)
    plt.title("Static Frequency")
    plt.plot(np.arange(N_iters) + 1, w_static, marker = "o")
    plt.xlabel("Lanczos step")
    plt.ylabel("Static frequency [cm-1]")
    plt.tight_layout()


    plt.figure(dpi = 160)
    plt.title("Green function without terminator")
    plt.imshow(dynamical_noterm, aspect = "auto", origin = "lower", extent = [W_START*CC.Phonons.RY_TO_CM, W_END*CC.Phonons.RY_TO_CM, 1, N_iters])
    plt.xlabel("Frequency [cm-1]")
    plt.ylabel("Steps")
    plt.colorbar()
    plt.tight_layout()


    plt.figure(dpi = 160)
    plt.title("Green function with terminator")
    plt.imshow(dynamical, aspect = "auto", origin = "lower", extent = [W_START*CC.Phonons.RY_TO_CM, W_END*CC.Phonons.RY_TO_CM, 1, N_iters])
    plt.xlabel("Frequency [cm-1]")
    plt.ylabel("Steps")
    plt.colorbar()
    plt.tight_layout()



    # get colormap
    cmap=plt.cm.viridis
    # build cycler with 5 equally spaced colors from that colormap
    c = cycler('color', cmap(np.linspace(0,1,N_iters)) )
    # supply cycler to the rcParam
    plt.rcParams["axes.prop_cycle"] = c

    plt.figure(dpi = 160)
    plt.title("Green function with terminator")
    for i in range(N_iters):
        plt.plot(w_array * CC.Phonons.RY_TO_CM, dynamical[i, :])

    # The last one is plotted in red underlined
    plt.plot(w_array * CC.Phonons.RY_TO_CM, dynamical[-1, :], color = "r", linewidth = 3.5)
    plt.xlabel("Frequency [cm-1]")
    plt.ylabel("Steps")
    plt.tight_layout()

    plt.show()

