#!/usr/bin/env python

import os
import sys
import json
import argparse
import numpy as np
import matplotlib
import multiprocessing
import logging

matplotlib.use('agg')
import matplotlib.pyplot as plt
from pychemia.code.abinit import AbinitInput, AbinitOutput
from pychemia.population.orbitaldftu import dmatpawu2params, params2dmatpawu, OrbitalDFTU
from pychemia.population.orbitaldftu import get_final_abinit_out
from pychemia.searcher import FireFly
from pychemia.db import get_database


def compare_params(path):
    if not os.path.isfile(path + os.sep + 'abinit.in'):
        print('ERROR: No abinit.in found at %s' % path)
        return
    # For making easier to see the values
    np.set_printoptions(linewidth=200, suppress=True)

    # Reading the INPUT
    abi = AbinitInput(path + os.sep + 'abinit.in')

    if 'lpawu' not in abi.variables:
        raise ValueError("Variable lpawu not found")
    ndim = 2 * max(abi['lpawu']) + 1

    idmatpawu = np.array(abi['dmatpawu']).reshape(-1, ndim, ndim)
    iparams = dmatpawu2params(idmatpawu, ndim)

    # Reading the OUTPUT
    abinitout = get_final_abinit_out(path)
    abo = AbinitOutput(abinitout)
    dmatpawu = abo.get_final_dmatpawu()
    odmatpawu = np.array(dmatpawu).reshape(-1, ndim, ndim)
    oparams = dmatpawu2params(odmatpawu, ndim)

    print('DMATPAWU')
    print('input')
    print(idmatpawu)
    print('output')
    print(odmatpawu)

    print('PARAMETRIC REPRESENTATION found at %s' % abinitout)
    for i in sorted(list(iparams.keys())):
        print(i)
        print('input')
        print(iparams[i])
        print('output')
        print(i)
        print(oparams[i])

    abo = AbinitOutput(abinitout)
    if not abo.is_finished:
        print('This output is not finished')
    try:
        nres2 = abo.get_energetics()['nres2'][-1]
        etot = abo.get_energetics()['etot'][-1]
        nscf = len(abo.get_energetics()['etot'])
        print("%30s ETOT: %15.6f NRES2: %15.6e NumSCF: %3d" % (path, etot, nres2, nscf))
    except:
        print("ERROR: Could not get energetics from %s" % abinitout)


def check_status(basepath):
    dirs = [x for x in os.listdir(basepath)
            if os.path.isdir(basepath + os.sep + x) and os.path.isfile(basepath + os.sep + x + os.sep + 'abinit.in')]
    print("%-40s %15s %15s %4s" % ("ABINIT output", "ETOT", 'nres2', 'nSCF'))
    for i in dirs:
        path = basepath + os.sep + i
        abinitout = get_final_abinit_out(path)
        if abinitout is None:
            continue

        abo = AbinitOutput(abinitout)
        if not abo.is_finished:
            continue
        try:
            nres2 = abo.get_energetics()['nres2'][-1]
            etot = abo.get_energetics()['etot'][-1]
            nscf = len(abo.get_energetics()['etot'])
            print("%-40s %15.6f %15.6e %4d" % (abinitout, etot, nres2, nscf))
        except:
            print("ERROR: Could not get final energetics from %s" % abinitout)


def plot_polar(popu, basepath):
    print('Plotting Euler Angles...')

    dirs = [x for x in os.listdir(basepath)
            if os.path.isdir(basepath + os.sep + x) and os.path.isfile(basepath + os.sep + x + os.sep + 'abinit.in')]
    fig = plt.figure(figsize=(21, 1.2 * len(dirs)))
    plt.subplots_adjust(left=0.0, bottom=0.0, right=0.99, top=0.99, wspace=None, hspace=None)
    etots = []

    for idir in dirs:
        pm, etot = popu.get_final_properties(basepath + os.sep + idir)
        etots.append(etot)

    etots = np.array(etots)