import os
import inspect
import warnings
import copy
import logging

from beluga.codegen.codegen import *

import dill
import numpy as np
import collections as cl

from beluga import problem, helpers
from beluga.bvpsol import algorithms, Solution
from .utils import tic, toc
from .utils.keyboard import keyboard
# from beluga.optimlib import ocp_to_bvp

config = dict(logfile='beluga.log',
              default_bvp_solver='Shooting',
              output_file='data.dill')

BVP = cl.namedtuple('BVP', 'deriv_func bc_func compute_control')
def bvp_algorithm(algo, **kwargs):
    """
    Helper method to load algorithm by name
    """
    # Load algorithm from the package
    for name, obj in inspect.getmembers(algorithms):
        if inspect.isclass(obj):
            if name.lower() == algo.lower():
                return obj(**kwargs)
    else:
        # Raise exception if the loop completes without finding an algorithm
        # by the given name
        raise ValueError('Algorithm '+algo+' not found')


def guess_generator(*args, **kwargs):
    """Helper for creating Initial guess generator."""
    guess_gen = problem.GuessGenerator()
    guess_gen.setup(*args,**kwargs)
    return guess_gen


def setup_beluga(logging_level=logging.INFO, display_level=logging.INFO, output_file=None):
    """Performs initial configuration on beluga."""

    # Get reference to the input file module
    frm = inspect.stack()[1]
    input_module = (inspect.getmodule(frm[0]))

    # Get information about input file
    info = inspect.getframeinfo(frm[0])

    # Suppress warnings
    warnings.filterwarnings("ignore")

    # Initialize logging system
    helpers.init_logging(logging_level,display_level, config['logfile'])

    # Set the output file name
    if output_file is not None:
        config['output_file'] = output_file

    # solve(problem)
    return


def solve(ocp, method, bvp_algorithm, steps, guess_generator, output_file='data.dill'):
    """
    Solves the OCP using specified method
    """

    # Initialize necessary conditions of optimality object
    # print("Computing the necessary conditions of optimality")
    logging.info("Computing the necessary conditions of optimality")
    from beluga.optimlib import methods

    wf = methods[method]
    # ocp_ws = ocp_to_bvp(ocp, guess_generator)
    ocp_ws = wf({'problem': ocp, 'guess': guess_generator})

    solinit = Solution()

    solinit.aux['const'] = dict((str(const.name),float(const.value))
                                for const in ocp_ws['constants'])
    solinit.aux['parameters'] = ocp_ws['problem_data']['parameter_list']

    # For path constraints
    solinit.aux['constraint'] = cl.defaultdict(float)
    solinit.aux['constraints'] = dict((s['name'], {'unit':str(s['unit']),
                                                   'expr':str(s['expr']),
                                                   'direction': s['direction'],
                                                   'arc_type': i,
                                                   'pi_list':[str(_) for _ in s['pi_list']]})
                                      for i, s in enumerate(ocp_ws['problem_data']['s_list'],1))
    solinit.aux['arc_seq'] = (0,)
    solinit.aux['pi_seq'] = (None,)
    bvp_fn, bvp = bvp_algorithm.preprocess(ocp_ws['problem_data'])
    # The initial guess is automatically stored in the bvp object
    # solinit is just a reference to it
    solinit = ocp_ws['guess'].generate(bvp_fn, solinit)

    # # includes costates
    state_names = ocp_ws['problem_data']['state_list']

    initial_states = solinit.y[:,0] # First column
    terminal_states = solinit.y[:,-1] # Last column
    initial_bc = dict(zip(state_names,initial_states))
    terminal_bc = dict(zip(state_names,terminal_states))

    solinit.aux['initial'] = initial_bc
    solinit.aux['terminal'] = terminal_bc
    tic()
    # TODO: Start from specific step for restart capability
    # TODO: Make class to store result from continuation set?
    out = dict()

    out['problem_data'] = ocp_ws['problem_data']

    ocp._scaling.initialize(ocp_ws)
    ocp_ws['scaling'] = ocp._scaling

    out['solution'] = run_continuation_set(ocp_ws, bvp_algorithm, steps, bvp_fn, solinit, bvp)
    total_time = toc()

    # tic()
    # out['solution'] = run_continuation_set(ocp_ws, bvp_algorithm, steps, bvp_fn, solinit2)
    # total_time = toc()

    logging.info('Continuation process completed in %0.4f seconds.\n' % total_time)
    bvp_algorithm.close()

    # Final time is appended as a parameter, so scale the output x variables to show the correct time
    for continuation_set in out['solution']:
        for sol in continuation_set:
            tf_ind = [i for i, s in enumerate(out['problem_data']['parameter_list']) if s is 'tf'][0]
            tf = sol.parameters[tf_ind]
            sol.t = sol.t*tf

    # Save data
    # del out['problem_data']['s_list']
    del out['problem_data']['states']
    del out['problem_data']['costates']

    qvars = out['problem_data']['quantity_vars']
    qvars = {str(k):str(v) for k,v in qvars.items()}
    out['problem_data']['quantity_vars'] = qvars
    with open(output_file, 'wb') as outfile:
        dill.settings['recurse'] = True
        dill.dump(out, outfile) # Dill Beluga object only

    return out['solution'][-1][-1]


def run_continuation_set(ocp_ws, bvp_algo, steps, bvp_fn, solinit, bvp):
    # Loop through all the continuation steps
    solution_set = []
    # Initialize scaling
    s = ocp_ws['scaling']
    problem_data = ocp_ws['problem_data']
    try:
        sol_guess = solinit
        sol = None
        for step_idx, step in enumerate(steps):
            logging.info('\nRunning Continuation Step #'+str(step_idx+1)+' : ')
            solution_set.append([])
            # Assign solution from last continuation set
            step.reset()
            step.init(sol_guess, problem_data)

            for aux in step:  # Continuation step returns 'aux' dictionary
                sol_guess.aux = aux

                logging.info('Starting iteration '+str(step.ctr)+'/'+str(step.num_cases()))
                tic()

                s.compute_scaling(sol_guess)
                sol_guess = s.scale(sol_guess)

                sol = bvp_algo.solve(bvp.deriv_func, None, bvp.bc_func, sol_guess)
                step.last_sol.converged = sol.converged
                sol = s.unscale(sol)

                if sol.converged:
                    # Post-processing phase

                    # Compute control history
                    # Required for plotting to work with control variables
                    sol.ctrl_expr = problem_data['control_options']
                    sol.ctrl_vars = problem_data['control_list']

                    #TODO: Make control computation more efficient
                    # for i in range(len(sol.x)):
                    #     _u = bvp.control_func(sol.x[i],sol.y[:,i],sol.parameters,sol.aux)
                    #     sol.u[:,i] = _u

                    ## DAE mode
                    # sol.u = sol.y[problem_data['num_states']:,:]
                    # Non-DAE:
                    f = lambda _t, _X: bvp_fn.compute_control(_t,_X,sol.parameters,sol.aux)
                    sol.u = np.array(list(map(f, sol.t, list(sol.y.T)))).T

                    # Copy solution object for storage and reuse `sol` in next
                    # iteration
                    solution_set[step_idx].append(copy.deepcopy(sol))
                    sol_guess = copy.deepcopy(sol)
                    elapsed_time = toc()
                    logging.info('Iteration %d/%d converged in %0.4f seconds\n' % (step.ctr, step.num_cases(), elapsed_time))
                else:
                    elapsed_time = toc()
                    logging.info('Iteration %d/%d failed to converge!\n' % (step.ctr, step.num_cases()))
    except Exception as e:
        import traceback
        traceback.print_exc()
        logging.error('Exception : ' + str(e))
        logging.error('Stopping')

    return solution_set
