#!/usr/bin/env python
# coding: utf-8
#
# Copyright © 2022 - Rennes Physics Institute
#
# This file is part of MsSpec-DFM.
#
# MsSpec-DFM is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# MsSpec-DFM is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# You should have received a copy of the GNU General Public License
# along with MsSpec-DFM.  If not, see <http://www.gnu.org/licenses/>.
#
# Source file  : src/python/msspec_dfm/cli.py
# Last modified: Fri, 25 Feb 2022 17:27:32 +0100
# Committed by : Sylvain Tricot <sylvain.tricot@univ-rennes1.fr> 1645806435 +0100

import click
import logging
import os
from msspec_dfm.input_file import InputFile
from msspec_dfm.version import __version__
from msspec_dfm.eps import epsilon
from msspec_dfm import plotting

@click.group(invoke_without_command=True, no_args_is_help=True)
@click.option('--version', is_flag=True, 
              help='Print the program version and exit.')
@click.option('-v', '--verbose', count=True,
              help='Increase the verbosity output. Full verbosity for -vvv')
def main(version, verbose):
    # Adjust the verbosity level
    verbosity_levels = [logging.ERROR, logging.WARNING, logging.INFO, logging.DEBUG]
    logging.basicConfig(level=verbosity_levels[min(verbose, len(verbosity_levels)-1)])

    # Print the version and exit
    if version:
        logging.debug('Getting the version number...')
        print(__version__)
        exit(0)


@main.command()
@click.option('--input-file', default='Data/epsilon.dat', show_default=True,
              metavar='FILE_NAME',
              help='The input filename.')
def generate(input_file):
    """Generate an input file.

       An input file with default parameters will be written to INPUT_FILE.
    """
    logging.debug("Generating the input file \'{:s}\'...".format(input_file))
    input_data = InputFile(input_file)
    input_data.write()
    logging.info("Input file with default values generated in: {}".format(input_file))


@main.command()
@click.argument('input-file')
@click.option('--folder', default='./Results', show_default=True,
              metavar='FOLDER_NAME',
              help='The output folder.')
def compute(input_file, folder):
    """Compute the dielectric function.

       All the calculation parameters are defined in INPUT_FILE.
       See help of the 'generate' command for a full reference of the file
       format and all the options.
    """
    logging.info("Computing...")
    logging.debug("    Input file   : \'{:s}\'".format(input_file))
    logging.debug("    Output folder: \'{:s}\'".format(folder))
    epsilon(input_file, folder)


@main.command()
@click.argument('plot-type', metavar='TYPE',
        type=click.Choice(['diel_func', 'susc_func', 'cond_func',
                           'loca_fiel', 'stru_fact', 'vertex_fu',
                           'deri_epsi', 'self_ener', 'dyna_coul'],
                           case_sensitive=False))
@click.option('--folder', default='./Results', show_default=True, metavar='FOLDER_NAME',
              help='The folder containing the results.')
@click.option('--img', is_flag=True,
              help='Select imaginery part.')
@click.option('--bounds', nargs=2, type=float, default=[-2, 2], show_default=True,
              metavar='LOWER UPPER',
              help='Limits for the colorbar scale.')
@click.option('--plot3d', is_flag=True,
              help='Plot in 3D.')
@click.option('--contour', is_flag=True,
              help='Add contours.')
@click.option('--pdeh', is_flag=True,
              help='Add the dispersion and electron-hole curves.')
@click.option('--levels', type=int, default=8, show_default=True,
              help='Number of levels for contour plots.')
@click.option('--Ef', 'Ef', type=float, default=1, show_default=True,
              help='Fermi energy if renormalization is needed.')
@click.option('--strides', type=int, default=64, show_default=True,
              help='Number of strides for 3D plot. The lower, the higher the resolution.')
def plot(plot_type, folder, img, bounds, plot3d, contour, pdeh, levels, Ef, strides):
    """Plotting the results.

       \b
       The plot TYPE value may be one of the following values:
         * diel_func
         * susc_func
         * cond_func
         * loca_fiel
         * stru_fact
         * vertex_fu
         * deri_epsi
         * self_ener
         * dyna_coul
    """
    logging.debug("Plotting...")
    for key, value in locals().items():
        logging.debug("    {}: {}".format(key, value))
    
    # data = plotting.Data(os.path.join(folder, plot_type + ".dat"))
    # data.load(imag=img)

    # pw = plotting.PlotWindow3d(data)
    # pw.plot()
    # pw.show()

    # start by loading data
    df_file = os.path.join(folder, 'diel_func.dat')
    pd_file = os.path.join(folder, 'plas_disp.dat')
    eh_file = os.path.join(folder, 'elec_hole.dat')
    X, Y, ReZ, ImZ, pd, eh = plotting.load_data(df_file, pd_file, eh_file)

    vmin, vmax = bounds
    Z = ReZ
    if img:
        Z = ImZ

    fig = ax = None

    if not(plot3d):
        fig, ax = plotting.plot_2d(X, Y, Z, cmap=plotting.cmap, vmin=vmin, vmax=vmax, fig=fig, ax=ax)
    else:
        fig, ax = plotting.plot_3d(X, Y, Z, cmap=plotting.cmap, vmin=vmin, vmax=vmax, fig=fig, ax=ax, strides=strides)

    if contour:
        fig, ax = plotting.plot_contour(X, Y, ReZ, levels=levels, vmin=vmin, vmax=vmax,
                            cmap=None, colors='black', fig=fig, ax=ax)

    if pdeh:
        fig, ax = plotting.plot_pdeh(pd, eh, Ef=Ef, fig=fig, ax=ax)

    plotting.entitle(fig, ax, img=img)
    plotting.plt.show()




if __name__ == '__main__':
    main()
