#!/usr/bin/env python
# coding: utf-8
"""This script runs RESOLVE.
    1. Loading the model formulation
    2. Loading data
    3. Creating a problem instance
    4. Solving the problem instance
    5. Loading results
    6. Writing results files
"""

# ############################ LICENSE INFORMATION ############################
# This file is part of the E3 RESOLVE Model.

# Copyright (C) 2019 Energy and Environmental Economics, Inc.
# For contact information, go to www.ethree.com

# The E3 RESOLVE Model is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# The E3 RESOLVE Model 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 Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with the E3 RESOLVE Model (in the file LICENSE.TXT). If not,
# see <http://www.gnu.org/licenses/>.
# #############################################################################

import ast
import os
import pandas as pd
import shutil
import sys

import pyomo.environ as pyo
from pyomo.opt import SolverFactory, TerminationCondition
from pyutilib.services import TempfileManager

# Resolve modules
import create_results_summary
import export_results
import fileio
import load_data
import model_formulation
import util


def create_problem_instance(scenario_inputs_directory, feature_toggles):
    """Load model formulation and data, and create problem instance.

    Args:
        scenario_inputs_directory (str): Path to scenario input data.
        feature_toggles (dict): Feature toggls.

    Returns:
        pyo.AbstractModel
    """
    # Get model, load data, and solve
    print('Reading model...')
    resolve_model = model_formulation.ResolveModel(feature_toggles)
    print('...model read.')

    if feature_toggles['deliverability']:
        print("NOTE: Deliverability functionality will be deprecated in future versions of RESOLVE.\
               \nDeliverability functionality can be replicated by splitting resource into:\
               \n1) With zero NQC (undelivered portion)\
               \n2) With non-zero NQC (delivered portion)")

    print('Loading data...')
    data, flexible_params = load_data.scenario_data(scenario_inputs_directory, resolve_model)
    print('...data read.')

    print('Compiling instance...')
    instance = resolve_model.create_instance(data)
    print('...instance created.')

    # example code for debugging via printing output
    # instance.THERMAL_RESOURCES.pprint()

    # just list out all the parameters that have changed value via the flexible_params functionality
    # so that people can look back at the log file for confirmation that the instance was created correctly
    if feature_toggles['debug']:
        print('\nNumber of values updated from Param default value via flexible params:')
        params_updated = flexible_params.index.get_level_values(level=0).unique()
        for param_name in params_updated:
            print('  {param_name}:'.format(param_name=param_name))

            model_components = flexible_params.loc[
                pd.IndexSlice[param_name, :]
            ].index.get_level_values(level=0).unique()

            for model_component in model_components:
                num_params_updated = len([
                    idx
                    for idx in getattr(instance, param_name)
                    if model_component in idx and
                    getattr(instance, param_name)[idx] != getattr(instance, param_name)._default_val
                ])
                if num_params_updated > 0:
                    print(
                        '    {model_component}: {num_params_updated}'
                        .format(model_component=model_component, num_params_updated=num_params_updated)
                    )

    return instance


def solve(instance, directory_structure, ilp_file_path=''):
    """Select solver and associated solver options.

    Args:
        instance (AbstractModel): Filled abstract model instance to solve.
        directory_structure (DirStructure): Directory structure and feature toggle object.
        ilp_file_path (str): File path to ILP file, which is where  Gurobi saves IIS (infeasible) model component

    Returns:
        solution (SolverResults): Solution to the abstract model instance.
    """

    # set cbc as the default solver, but allow for other command line arguments
    if len(sys.argv) >= 3:
        solver_name = sys.argv[2]
    else:
        solver_name = "cbc"

    # TODO: Use a proper command line args package (like argparse)
    # the only time this should be 4 is if we're using Gurobi cloud interface (check cloud API for available job slots)
    if 'cloud' in sys.argv and solver_name == 'gurobi':
        cloud = True
    else:
        cloud = False

    # ### Solve ### #
    solver_io = 'lp'
    # use the AMPL interface for XPRESS to get IIS info out of the model
    if solver_name == 'amplxpress':
        solver_io = 'nl'
    # use the local CBC executable by default
    if solver_name == 'cbc':

        import platform
        cur_sys = platform.system()

        if cur_sys == "Darwin":
            cur_sys = "macOS"

        print("Current platform is {}".format(cur_sys))

        if cur_sys == "Windows":
            cbc_path = "cbc.exe"
        else:
            cbc_path = "cbc"

        solver = SolverFactory(
            solver_name,
            executable=os.path.join(directory_structure.DIRECTORY,
                                    'solvers',
                                    cbc_path),
            solver_io=solver_io
        )
    else:
        solver = SolverFactory(solver_name, solver_io=solver_io)

    # Hard-coded solver options not defined in solver_options.csv file
    if solver_name == "gurobi":
        solver.options["ResultFile"] = ilp_file_path
    elif solver_name == "amplxpress":
        solver.options["IIS"] = ""
        # Create an 'iis' (Irreducible Infeasible Set) suffix component on the instance.
        # If the solver supports suffixes, Pyomo will receive the stored information.
        # This should work for any AMPL-interfaced solver, but the only one we use that way
        # currently is XPRESS.
        instance.iis = pyo.Suffix(direction=pyo.Suffix.IMPORT)

    # Read in user-defined solver options
    user_solver_options = pd.read_csv(
        os.path.join(directory_structure.SCENARIO_INPUTS_DIRECTORY, 'solver_options.csv'),
        header=0,
        index_col=[0, 1]
    )
    if solver_name in user_solver_options.index.get_level_values(0):
        for option_name, option in user_solver_options.loc[pd.IndexSlice[solver_name], :].iterrows():
            if 'bool' in option['dtype'].lower():
                solver.options[option_name] = ast.literal_eval(option['value'])
            elif 'int' in option['dtype'].lower():
                solver.options[option_name] = int(option['value'])
            elif 'float' in option['dtype'].lower():
                solver.options[option_name] = float(option['value'])
            else:
                solver.options[option_name] = option['value']

    # Optionally keep human-readable LP files for debugging
    if directory_structure.feature_toggles['debug']:
        keepfiles = True
        symbolic_solver_labels = True
    else:
        keepfiles = False
        symbolic_solver_labels = False

    if solver_name == "gurobi" and cloud:
        # only load get_gurobi_jobs if `cloud` is used
        try:
            import get_gurobi_jobs

            license_file = get_gurobi_jobs.check_cloud_status()
            os.environ['GRB_LICENSE_FILE'] = os.path.abspath(license_file)
        except ImportError as error:
            print(
                'ImportError: {}. If you do not need to use Gurobi Instant Cloud and do not have access set up, please do not use the options "gurobi cloud".'
                .format(error.message)
            )

    print('Starting solver...')
    solution = solver.solve(instance, keepfiles=keepfiles, tee=True, symbolic_solver_labels=symbolic_solver_labels)

    # escape if model is infeasible
    if solution.solver.termination_condition == TerminationCondition.infeasible:
        if solver_name == 'gurobi':
            raise Exception(
                'Model was infeasible. Gurobi has saved the file {} with possible cause of infeasibility.'.
                format(ilp_file_path)
            )
        elif solver_name == 'amplxpress':
            raise Exception(
                'Solver identified the following constraints to be infeasible: \n' +
                '\n  '.join(sorted(c.name for c in instance.iis))
            )
        else:
            raise Exception('Model was infeasible.')

    return solution


def get_objective_function_value(instance, scenario_results_directory):
    """Save the objective function value.

    Args:
        instance (pyo.AbstractModel): Filled abstract model instance to solve.
        scenario_results_directory ([type]): Scenario results directory to write to.
    """
    print('\nObjective function value is: {:,.2f}'.format(instance.Total_Cost()))
    with fileio.filewriter(
            os.path.join(scenario_results_directory, 'objective_function_value.txt')
        ) as objective_file:
        objective_file.write("Objective function value is: " + str(instance.Total_Cost()))


def run_scenario(directory_structure, ilp_file_path=''):
    """Run a scenario. Determine and create scenario directories, create problem instance, solve, and export results.

    Args:
        directory_structure (DirStructure): Directory structure and feature toggle object.
    """

    # Directories
    scenario_inputs_directory = os.path.join(directory_structure.SCENARIO_INPUTS_DIRECTORY)
    scenario_results_directory = os.path.join(directory_structure.SCENARIO_RESULTS_DIRECTORY)
    scenario_logs_directory = os.path.join(directory_structure.SCENARIO_LOGS_DIRECTORY)

    # Write logs to this directory
    TempfileManager.tempdir = scenario_logs_directory

    # Create problem instance
    instance = create_problem_instance(scenario_inputs_directory, directory_structure.feature_toggles)

    # Create a 'dual' suffix component on the instance, so the solver plugin will know which suffixes to collect
    instance.dual = pyo.Suffix(direction=pyo.Suffix.IMPORT)

    # Solve
    solution = solve(instance, directory_structure, ilp_file_path)

    # Get objective function value
    get_objective_function_value(instance, scenario_results_directory)

    # Load and export results
    export_results.export_results(instance,
                                  solution,
                                  scenario_results_directory,
                                  debug_mode=1)

    # Create results summaries
    create_results_summary.create_summaries(scenario_results_directory)

    # Copy inputs passthrough to summary results for Results Tool
    shutil.copyfile(
        scenario_inputs_directory + r'/inputs_passthrough.csv',
        scenario_results_directory + r'/summary/inputs_passthrough.csv'
    )
    shutil.copyfile(
        scenario_inputs_directory + r'/resource_results_grouping.csv',
        scenario_results_directory + r'/resource_results_grouping.csv'
    )

    print('Done.')


def main():
    """Main methods that calls all submethods."""

    # Scenario to run is given as first script argument (0th argument is the script name)
    # this must be the same name as a folder in the 'inputs' directory
    scenario_name = sys.argv[1]

    code_directory = os.getcwd()

    dir_str = util.DirStructure(code_directory, scenario_name)
    dir_str.make_directories()
    dir_str.get_feature_toggles()

    # Create a custom logger that prints warnings to scenario results directory
    new_logger = util.get_new_logger(dir_str)

    # Keep old logger functionality for now
    logger = util.Logger(dir_str)
    log_file = logger.log_file_path
    print('Running scenario {}...'.format(scenario_name))
    print('Logging run to {}...'.format(log_file))
    stdout = sys.stdout
    sys.stdout = logger
    run_scenario(dir_str, logger.ilp_file_path)
    sys.stdout = stdout  # return sys.stdout to original, just in case


if __name__ == "__main__":
    main()
