#! /usr/bin/env python
# coding=utf-8
""" Find best PSF stars using GA to minimize mean error.

    .. seealso:: :ref:`gapick`
"""
from __future__ import absolute_import, division, print_function
__metaclass__ = type

# For now:
# TODO: --neigbours option: once for generation neightbour removal
# TODO: --timestamp: add current timestamp to directory name (+)
# TODO: conf option for
# TODO: warnings about: daophot.opt, allstar.opt
# TODO: single parameter --aperture, instead of photo.opt
# TODO: write down text based statistics from logbook.pkl
# TODO: some docs, and link to them in --help
# TODO: allstar.opt missing in result dir!
# For later, on request:
# TODO: Option for provide own PSF stars set, used for comaprision, including in candidates, put in initial population
# TODO: Checkpoints - continue calculation

import os
import logging
import random
import pickle
import time
from datetime import timedelta
from copy import deepcopy

try:
    from itertools import izip_longest as zip_longest
except ImportError:
    from itertools import zip_longest

import numpy
from bitarray import bitarray
from scipy.stats import sigmaclip
from deap import base
from deap import creator
from deap import tools

import astwro.tools.__commons as commons
import astwro.starlist as sl
import astwro.pydaophot as dao
import astwro.tools
import astwro.utils as utils
from astwro.phot import PhotError

_time_format = '%a %H:%M:%S'

# Definitions for genetic algorithm:
# 'individual' - subset of candidates competing with other subsets to be the best in fitness
# 'genome' - the individual's definition with subset of candidates stars, is represented by binary string
#     of length equal to number of candidate stars. '1' means that star is in subset.
# 'fitness' - minimized function on individual, mean of errors form allstar if individual's stars are
#     taken as PSF stars
# 'population' - set of all individuals in some iteration

# 2.1 Definition of routines used by algorithms: initializations, scoring
def select_stars(starlist, genome):
    # type: (sl.StarList, bitarray) -> sl.StarList
    # Select stars present in genome
    return starlist[genome.tolist()]

def random_genome(ind_class, len, prob_limits):
    # type: (type, int, [float, float]) -> bitarray
    # Create random genome of length `len` in which probability of '1' on any position is from `prob_limits` range.
    prob = random.uniform(prob_limits[0], prob_limits[1])
    return ind_class([random.random() <= prob for _ in range(len)])

def clone_individual(individual):
    # type: (bitarray) -> bitarray
    # Make a deepcopy-clone - deepcopy genome bitarray and associated fitness
    n = deepcopy(individual)
    n.fitness = deepcopy(individual.fitness)
    return n

def calc_spectrum(pop):
    # calculates 'spectrogram' of population
    # which is a statistic of stars occurrences in population individuals
    spec = numpy.zeros(len(pop[0]))
    for ind in pop:
        spec += ind.tolist()
    return spec

# === astwro version <= 0.7.0
# def fitness_for_als(als):
#     # type: (sl.StarList) -> (float,)
#     # Calucalates fitness from allstar result
#     clipped = sigmaclip(als.chi)[0]
#     return numpy.sqrt((clipped * clipped).sum()),  # fitness is tuple (val,)

# === astwro version > 0.7.0
# def fitness_for_als(als):
#     # type: (sl.StarList) -> (float,)
#     # Calucalates fitness from allstar result
#     w = 100*(-als.mag/5)
#
#     return (als.chi*w).sum()/w.sum(),  # fitness is tuple (val,)

# === astwro version > 0.7.2
def fitness_for_als(als):
    # type: (sl.StarList) -> (float,)
    # Calucalates fitness from allstar result
    w = 100*(-als.mag/5)

    return (als.mag_err*w).sum()/w.sum(),  # fitness is tuple (val,)

def eval_population(population, candidates, workers, show_progress, fine_tune):
    if fine_tune:
        return eval_population_fine_psf(population, candidates, workers, show_progress)
    else:
        return eval_population_simple(population, candidates, workers, show_progress)

def eval_population_fine_psf(population, candidates, workers, show_progress):
    # type: (list(bitarray), sl.StarList, list(dict), bool) -> list
    # Evaluates fitness for all individual in population.

    # This version uses sofisticated process from daophot_bialkow
    # Uses daophots and allstars processes form `workers` running them in parallel.
    # :return: list fitnesses (1-element couples as `deap` lib likes)

    progress = None
    if show_progress:
        progress = utils.progressbar(total=len(population), step=len(workers))
        progress.print_progress(0)
    fitnesses = []
    f_max = None
    # https://docs.python.org/3/library/itertools.html#itertools-recipes grouper()
    # grouping population into chunks of size `parallel`
    for chunk in zip_longest(*([iter(population)] * len(workers))):
        active = [g is not None for g in chunk]  # all active, on errors some could be deactivated
        # PSF
        for individual, worker in zip(chunk, workers):
            if individual:
                pdf_s = select_stars(candidates, individual)
                worker['daophot'].write_starlist(pdf_s, 'i.lst')
                worker['daophot'].PSf(psf_stars='i.lst')  # add to queue only - batch mode
                worker['daophot'].run(wait=False)  # parallel start all workers without waiting
        # wait for PSF and start ALLSTAR nei
        for i, worker in enumerate(workers):
            if active[i]:
                worker['daophot'].wait_for_results()  # now wait before using results
                active[i] = worker['daophot'].PSf_result.converged  # PSF is not always successful
                if active[i]:
                    worker['allstar'].ALlstar(stars='i.nei')  # enqueue calculation
                    worker['allstar'].run(wait=False)  # asynchronous / parallel
        # second PSF
        for i, worker in enumerate(workers):
            if active[i]:
                worker['allstar'].wait_for_results()
                if worker['allstar'].ALlstars_result.success:
                    worker['daophot'].SUbstar(subtract='i.als', leave_in='i.lst')
                    worker['daophot'].run()  # quick run
                    worker['daophot'].ATtach('is')
                    worker['daophot'].PSf(photometry='i.als', psf_stars='i.lst')
                    worker['daophot'].run(wait=False)
                else:
                    active[i] = False
        # second allstar
        for i, worker in enumerate(workers):
            if active[i]:
                worker['daophot'].wait_for_results()
                if worker['daophot'].PSf_result.success:
                    worker['allstar'].ALlstar(stars='i.nei')  # enqueue calculation
                    worker['allstar'].run(wait=False)  # asynchronous / parallel
                else:
                    active[i] = False
        # third PSF
        for i, worker in enumerate(workers):
            if active[i]:
                worker['allstar'].wait_for_results()
                if worker['allstar'].ALlstars_result.success:
                    worker['daophot'].SUbstar(subtract='i.als', leave_in='i.lst')
                    worker['daophot'].run()  # quick run
                    worker['daophot'].ATtach('is')
                    worker['daophot'].PSf(photometry='i.als', psf_stars='i.lst')
                    worker['daophot'].run(wait=False)
                else:
                    active[i] = False
        # final allstar
        for i, worker in enumerate(workers):
            if active[i]:
                worker['daophot'].wait_for_results()
                if worker['daophot'].PSf_result.success:
                    worker['allstar'].ALlstar(stars='als.ap')  # enqueue calculation
                    worker['allstar'].run(wait=False)  # asynchronous / parallel
                else:
                    active[i] = False

        # wait for allstar for workers and process results
        for i, worker in enumerate(workers):
            if active[i]:
                worker['allstar'].wait_for_results()
                all_s = worker['allstar'].ALlstars_result.als_stars
                f = fitness_for_als(all_s)
                fitnesses.append(f)  # fitness is tuple (val,)
                f_max = f if f_max is None or f_max[0] < f[0] else f_max
            #                f_max = f[0] if f_max is None else max(f_max, f[0])
            else:
                fitnesses.append(None)
        # cut fitnesses if longer than population
        # (finesses mod parallel == 0, some None at the end can appear if population mod parallel != 0)
        fitnesses = fitnesses[:len(population)]
        # fill gaps in fitnesses with maximum of rest of population
        for i, f in enumerate(fitnesses):
            if f is None:
                fitnesses[i] = f_max
        if progress:
            progress.print_progress()

    return fitnesses

def eval_population_simple(population, candidates, workers, show_progress):
    # type: (list(bitarray), sl.StarList, list(dict), bool) -> list
    # Evaluates fitness for all individual in population.
    # Uses daophots and allstars processes form `workers` running them in parallel.
    # :return: list fitnesses (1-element couples as `deap` lib likes)

    progress = None
    if show_progress:
        progress = utils.progressbar(total=len(population), step=len(workers))
        progress.print_progress(0)
    fitnesses = []
    f_max = None
    # https://docs.python.org/3/library/itertools.html#itertools-recipes grouper()
    # grouping population into chunks of size `parallel`
    for chunk in zip_longest(*([iter(population)] * len(workers))):
        active = [g is not None for g in chunk]  # all active, on errors some could be deactivated
        # PSF
        for individual, worker in zip(chunk, workers):
            if individual:
                pdf_s = select_stars(candidates, individual)
                worker['daophot'].PSf(psf_stars=pdf_s)  # add to queue only - batch mode
                worker['daophot'].run(wait=False)  # parallel start all workers without waiting
        # wait for PSF and start ALLSTAR
        for i, worker in enumerate(workers):
            if active[i]:
                worker['daophot'].wait_for_results()  # now wait before using results
                active[i] = worker['daophot'].PSf_result.converged  # PSF is not always successful
                if active[i]:
                    worker['allstar'].ALlstar(stars='als.ap')  # enqueue calculation
                    worker['allstar'].run(wait=False)  # asynchronous / parallel
        # wait for allstar for workers and process results
        for i, worker in enumerate(workers):
            if active[i]:
                worker['allstar'].wait_for_results()
                all_s = worker['allstar'].ALlstars_result.als_stars
                f = fitness_for_als(all_s)
                fitnesses.append(f)  # fitness is tuple (val,)
                f_max = f if f_max is None or f_max[0] < f[0] else f_max
            #                f_max = max(f_max, f)
            else:
                fitnesses.append(None)
        # cut fitnesses if longer than population
        # (finesses mod parallel == 0, some None at the end can appear if population mod parallel != 0)
        fitnesses = fitnesses[:len(population)]
        # fill gaps in fitnesses by maximum of rest of population
        for i, f in enumerate(fitnesses):
            if f is None:
                fitnesses[i] = f_max
        if progress:
            progress.print_progress()

    return fitnesses

def _prepare_output_dir(outdir, overwrite, srcdir, arg):
    # type: (str, bool, str, object) -> (utils.CycleFile, utils.CycleFile, utils.CycleFile, str)
    # Prepare output directory for results
    if outdir:
        from shutil import copytree, rmtree

        outdir = os.path.abspath(os.path.expanduser(outdir))
        if os.path.exists(outdir):
            if overwrite:
                rmtree(outdir)
            else:
                logging.error('--out-dir:{} already exists and no --overwrite requested.'.format(outdir))
                raise Exception('Output directory {} exists.'.format(outdir))
        copytree(srcdir, outdir)
        logging.info('Results dir created: {}'.format(outdir))
        # todo allstar.opt missing in result dir!
        basename = 'gen'
        lst_file = utils.cyclefile(outdir, basename, '.lst')
        reg_file = utils.cyclefile(outdir, basename, '.reg')
        gen_file = utils.cyclefile(outdir, basename, '.gen')
        # write down parameters
        with open(os.path.join(outdir, 'about.txt'), 'w') as f:
            print("astwro pydaophot/tools version: {}/{}\n".format(dao.__version__, astwro.tools.__version__), file=f)
            for k, v in arg.__dict__.items():
                print('{}\t= {}'.format(k, v), file=f)
        return lst_file, reg_file, gen_file, outdir
    else:
        return None, None, None, None

def _plot_err(pe, pe_psf, psf, fig_file):
    # type: (PhotError) -> None
    try:
        import matplotlib
        matplotlib.use('pdf')
        import matplotlib.pyplot as plt
    except ImportError:
        logging.warning("'matplotlib' cannot be imported. No control diagrams will be generated")
    except:
        logging.warning('matplotlib PDF output cannot be initialized. No control diagrams will be generated')
        return
    else:
        try:
            plt.rc('text', usetex=True)
        except:
            pass

        ax = plt.subplot()
        ls = numpy.linspace(pe.mag.min(), pe.mag.max())
        ax.plot(pe.mag[~pe.outlayers_mask], pe.err[~pe.outlayers_mask], '.', alpha=0.7, ms=2,
                label='training set')
        ax.errorbar(pe.mag_means, pe.err_means, yerr=pe.err_means_stderr * 1.0, fmt='D', markerfacecolor='none', alpha=0.7, lw=1,
                    label='error bins')
        ax.plot(ls, pe.evaluate_fit(ls), lw=1,
                label='error fit')
        ax.plot(ls, pe.evaluate_fit_plus_sigmas(ls),  '--', lw=1,
                label='error$+{:.1f}\\sigma$ fit (training limit)'.format(pe.fitplussigma))
        ax.plot(ls, pe_psf.evaluate_fit_plus_sigmas(ls),  '--', lw=1,
                label='error$+{:.1f}\\sigma$ fit (PSF limit)'.format(pe_psf.fitplussigma))
        ax.plot(pe.mag[pe.outlayers_mask], pe.err[pe.outlayers_mask], '.', alpha=0.7, ms=3,
                label='outlayers')
        ax.plot(psf.mag, psf.mag_err, '.', alpha=0.7, ms=3,
                label='PSF candidates')
        ax.set_xlabel('mag'); ax.set_ylabel('mag error')
        ax.legend()
        ax.figure.savefig(fig_file, bbox_inches='tight')
        ax.set_ylim(0,0.2)
        ax.figure.savefig('{}_zoomed{}'.format(*os.path.splitext(fig_file)), bbox_inches='tight')
        logging.info('Control plot mag/err {} generated'.format(fig_file))

def __do(arg):
    # Main routine, common for command line, and python scripts call
    # :type arg: Namespace

    start_time = time.time()

    # Configure progress logging
    clogger = logging.getLogger('clean logger')  # create another logger for stats without prefixes (clean)
    chandler = logging.StreamHandler()
    chandler.setFormatter(logging.Formatter('%(message)s'))
    clogger.propagate = False
    clogger.addHandler(chandler)

    check_arguments(arg)

    # image_file
    if arg.image is None:
        from astwro.sampledata import fits_image
        arg.image = fits_image()  # sample image

        logging.warning('No image file argument provided (-h for help), using demonstration image: %s', arg.image)

    # expand patches of file parameters
    arg.all_stars_file = dao.Daophot.expand_path(arg.all_stars_file)
    arg.psf_stars_file = dao.Daophot.expand_path(arg.psf_stars_file)
    arg.photo_opt = dao.Daophot.expand_path(arg.photo_opt)
    arg.daophot_opt = dao.Daophot.expand_path(arg.daophot_opt)

    # get single daophot and ATtach file
    dp = dao.Daophot(image=arg.image, daophotopt=arg.daophot_opt)
    al = dao.Allstar(image=arg.image, dir=dp.dir)

    # all stars file
    if not arg.all_stars_file:
        logging.warning('No all-stars-file provided, Stars will be found by daophot FIND (frames av/sum: %d/%d)',
                        arg.frames_av, arg.frames_sum)
        find = dp.FInd(arg.frames_av, arg.frames_sum)
        logging.info('{} stars found by FIND, sky estimation: {}, err/dev: {}/{}'.format(find.stars, find.sky, find.err,
                                                                                find.skydev))
        arg.all_stars_file = find.starlist_file

    # photometry
    if not arg.photo_opt:  # no file no problem, but recreate default radius if not provided
        if arg.photo_is == 0:
            arg.photo_is = 35
        if arg.photo_os == 0:
            arg.photo_os = 50
        if not arg.photo_ap:
            arg.photo_ap = [8]
    logging.info('Performing preliminary {} pixels aperture  photometry'.format(arg.photo_ap[0]))
    photometry = dp.PHotometry(photoopt=arg.photo_opt,
                               IS=arg.photo_is,
                               OS=arg.photo_os,
                               apertures=arg.photo_ap,
                               stars=arg.all_stars_file)

    # pick PFS stars candidates
    PICKcandidates = not arg.psf_stars_file
    if PICKcandidates:
        pick = dp.PIck(arg.stars_to_pick, arg.faintest_to_pick)
        arg.psf_stars_file = pick.picked_stars_file

    # psf (for errors collection)
    PSf_result = dp.PSf(psf_stars=arg.psf_stars_file)

    # all stars (filtering by photometry errors and magnitudes)
    stars = photometry.photometry_starlist
    count0 = stars.shape[0]
    stars = stars[stars.mag < arg.max_ph_mag]
    count1 = stars.shape[0]
    stars = stars[stars.mag_err < arg.max_ph_err]
    count2 = stars.shape[0]
    logging.info('{} (of {}) stars left after filtering against {:.2f} photometry magnitude threshold, '
                 .format(count1, count0, arg.max_ph_mag))
    logging.info('{}  left after {:.2f} aperture photometry error threshold'
                 .format(count2, arg.max_ph_err))

    logging.info('Performing preliminary PSF profile photometry')
    ALlstar_result = al.ALlstar(stars=stars)
    logging.info('{} stars converged in PSF profile photometry'
                 .format(ALlstar_result.stars_no[0]))
    logging.info('Fitting PSF profile photometry errors')
    pe = PhotError(ALlstar_result.als_stars.mag, ALlstar_result.als_stars.mag_err, fitplussigma=arg.clip_fit_sigmas)
    logging.info('{} stars left, after excluding profile photometry error {} sigma outliers'
                 .format((~pe.outlayers_mask).sum(), arg.clip_fit_sigmas))
    stars = ALlstar_result.als_stars[~pe.outlayers_mask]

    # for the sake of optimisation we write starlist to file to avoid multiple serialization of this list
    # (instead of more obvious providing starlist as the allstar argument)
    dp.write_starlist(stars, 'als.ap', sl.DAO.AP_FILE)

    # candidates (filter out big psf errors)
    pe2 = None
    if PICKcandidates: # do it again
        if arg.clip_fit_sigmas > arg.clip_psf_fit_sigmas: # tighter filtering for PSF candids
            logging.info('Excluding {} sigma outlayers for PICK'.format(arg.clip_psf_fit_sigmas))
            pe2 = PhotError(ALlstar_result.als_stars.mag, ALlstar_result.als_stars.mag_err,
                           fitplussigma=arg.clip_psf_fit_sigmas)
            cands_for_cands = ALlstar_result.als_stars[~pe2.outlayers_mask]
        else:
            cands_for_cands = stars
        pick = dp.PIck(arg.stars_to_pick, arg.faintest_to_pick, photometry=cands_for_cands)
        logging.info('{} PSF stars candidates selected by PICK'.format(pick.stars))
        arg.psf_stars_file = pick.picked_stars_file
        PSf_result = dp.PSf(psf_stars=arg.psf_stars_file) # collect PSF errors
    candidates = dp.read_starlist(arg.psf_stars_file, add_psf_errors=True)
    org_cand_no = candidates.stars_number()
    err = PSf_result.errors
    averr = err.psf_err.mean()
    err = err[(err.psf_err < arg.max_psf_err_mult * averr) & (
                err.flag == ' ')]  # filter out big errors and * or ? marked stars
    candidates = candidates.loc[err.index]

    # candidates = candidates[candidates.psf_err < arg.max_psf_err]  # old filter, new above
    logging.info(
        "{} good candidates ({} rejected: * or ? or psf error exceeded max-psf-err-mult*averge-error "
        "= {:.2f}*{:.2f} = {:.2f})"
            .format(
            candidates.stars_number(),
            org_cand_no - candidates.stars_number(),
            arg.max_psf_err_mult,
            averr,
            arg.max_psf_err_mult * averr)
    )
    if arg.include_psf:
        logging.info('PSF star candidates stays in training set (--include-psf)')
    else:
        overlapping = stars.index.isin(candidates.index)
        stars = stars[~overlapping]
        logging.info('{} stars left in training set after exclusion PSF star candidates (no --include-psf)'
                     .format(stars.stars_number()))

    if candidates.stars_number() < 15:
        logging.error("Number of candidates lass than 15. GA needs more. Sorry...")

    if stars.stars_number() < 15:
        logging.error("Number stars in training set lass than 15. GA needs more. Sorry...")

    # Prepare output directory for results
    lst_file, reg_file, gen_file, result_dir = _prepare_output_dir(arg.out_dir, arg.overwrite, str(dp.dir), arg)
    _plot_err(pe, pe2, candidates, os.path.join(arg.out_dir, 'prel_psf_filtering.pdf'))

    #  From all candidates find best subset, where best means minimizing mean of errors form allstar
    #       using genetic algorithm
    #  For details about setting up genetic algorithm with DEAP visit:
    #       http://deap.gel.ulaval.ca/doc/default/examples/ga_onemax.html

    creator.create("FitnessMax", base.Fitness, weights=(-1.0,))
    creator.create("Individual", bitarray, fitness=creator.FitnessMax)

    toolbox = base.Toolbox()
    # Structure initializes
    toolbox.register("individual", random_genome, creator.Individual, candidates.stars_number(), arg.ga_init_prob)
    toolbox.register("population", tools.initRepeat, list, toolbox.individual)
    toolbox.register('clone', clone_individual)

    # The Genetic Operators

    # set min_stars to all_cand_no*ga_init_prob/2
    toolbox.register('mate', tools.cxTwoPoint)
    toolbox.register('mutate', tools.mutFlipBit, indpb=arg.ga_mut_str)
    toolbox.register('select', tools.selTournament, tournsize=3)

    # setup stats
    stats_fits = tools.Statistics(key=lambda ind: ind.fitness.values)
    stats_star = tools.Statistics(key=sum)  # number of stars is an sum of bitarray: [001101010001] has 5 stars
    stats = tools.MultiStatistics(fitness=stats_fits, size=stats_star)
    stats.register('avg', numpy.mean)
    stats.register('std', numpy.std)
    stats.register('min', numpy.min)
    stats.register('max', numpy.max)

    # Initiate workers. Each worker has Daophot and Allstar objects sharing runner directory,
    # working in batch mode.
    workers = []
    workers_logger = logging.getLogger('worker')
    workers_logger.setLevel('ERROR')  # prevent workers flood output with logrecords
    for i in range(arg.parallel):
        d = dp.clone()  # clone previously used daophot
        d.batch_mode = True
        a = dao.Allstar(dir=d.dir, image=d.image, batch=True, options={'MA': 100})
        d.logger = workers_logger
        a.logger = workers_logger
        workers.append({'daophot': d, 'allstar': a})

    # Setup initial population, HoF and logbook and  or load it from checkpoint when continuing previous calculation
    start_gen = 0

    #    if arg.checkpoint:   ## Not implemented
    if False:
        with open(os.path.expanduser(arg.checkpoint), "rb") as f:
            checkpoint = pickle.load(f)
        pop = checkpoint['population']
        start_gen = checkpoint['generation']
        hof = checkpoint['halloffame']
        logbook = checkpoint['logbook']
        logging.info('Restoring genetic algorithm on {} of {} generations'.format(start_gen, arg.ga_max_iter))
    else:
        hof = tools.HallOfFame(maxsize=10)

        logbook = tools.Logbook()
        logbook.header = 'gen', 'fitness', 'size'
        logbook.chapters['fitness'].header = 'min', 'avg', 'max', 'std'
        logbook.chapters['size'].header = 'min', 'avg', 'max'

        pop = toolbox.population(n=arg.ga_pop)
        logging.info('Starting genetic algorithm for {} generations at {}'.format(
            arg.ga_max_iter,
            time.strftime(_time_format, time.localtime())
        ))
        logging.info('{} parallel threads, ETA will be calculated after generation 1'.format(arg.parallel))

    # Calculate fitnesses of initial population
    fitnesses = eval_population(pop, candidates, workers, show_progress=not arg.no_progress, fine_tune=arg.fine)
    for ind, fit in zip(pop, fitnesses):
        ind.fitness.values = fit

    record = stats.compile(pop)
    logbook.record(gen=0, spectrum=calc_spectrum(pop), **record)
    clogger.info('{}\t ETA: [... to be determined]'.format(logbook.stream))

    evolution_start_time = time.time()

    # Begin the evolution
    for g in range(start_gen + 1, arg.ga_max_iter):
        #  New Generation
        #  select the next generation individuals
        offspring = toolbox.select(pop, len(pop))
        # Clone the selected individuals
        offspring = list(map(toolbox.clone, offspring))
        # Apply crossover and mutation on the offspring
        for child1, child2 in zip(offspring[::2], offspring[1::2]):
            if random.random() < arg.ga_cross_prob:
                toolbox.mate(child1, child2)
                del child1.fitness.values
                del child2.fitness.values
        for mutant in offspring:
            if random.random() < arg.ga_mut_prob:
                toolbox.mutate(mutant)
                del mutant.fitness.values

        # calculate fitnesses of new individuals
        invalid_ind = [ind for ind in offspring if not ind.fitness.valid]
        fitnesses = eval_population(invalid_ind, candidates, workers, show_progress=not arg.no_progress,
                                    fine_tune=arg.fine)
        for ind, fit in zip(invalid_ind, fitnesses):
            ind.fitness.values = fit
        # New population from offspring
        pop[:] = offspring

        # Stats
        # hof.update(pop)  # not implemented yet, __deapcopy__ of the Individual should work first
        ETA = time.strftime(_time_format, time.localtime(
            evolution_start_time + (time.time() - evolution_start_time) * arg.ga_max_iter / g))
        record = stats.compile(pop)
        logbook.record(gen=g, spectrum=calc_spectrum(pop), **record)
        clogger.info('{}\t ETA: {}'.format(logbook.stream, ETA))

        # For every generation create lst file and ds9 reg file of best and point symlinks to last generation
        if result_dir:
            best_ind = tools.selBest(pop, 1)[0]
            best_stars = select_stars(candidates, best_ind)
            lst_file.next_file(g)
            reg_file.next_file(g)
            gen_file.next_file(g)
            sl.write_dao_file(best_stars, lst_file.file, sl.DAO.LST_FILE)
            sl.write_ds9_regions(best_stars, reg_file.file)
            for ind in pop:
                gen_file.file.write(ind.to01() + '\n')
            with open(os.path.join(result_dir, 'logbook.pkl'), 'wb') as f:
                pickle.dump(logbook, f)
            checkpoint = dict(population=pop, generation=g, halloffame=hof, logbook=logbook)
            with open(os.path.join(result_dir, 'checkpoint.chk'), 'wb') as f:
                pickle.dump(checkpoint, f)

        # end of evolution loop

    if result_dir:
        lst_file.close()
        reg_file.close()
        gen_file.close()

    logging.info('Successful evolution finished at {} (elapsed time: {:s})'.format(
        time.strftime(_time_format, time.localtime()),
        str(timedelta(seconds=time.time() - start_time))
    ))

    best_ind = tools.selBest(pop, 1)[0]
    logging.info('Best individual is {}, {}'.format(best_ind, best_ind.fitness.values))

    best_stars = select_stars(candidates, best_ind)
    return best_stars

__arg_parser_singlethon = None

def __arg_parser():
    global __arg_parser_singlethon
    if __arg_parser_singlethon:
        return __arg_parser_singlethon
    import argparse
    parser = argparse.ArgumentParser(
        description='Find best PSF stars using genetic algorithm running the daophot/allstar'
                    ' workflow for number of a PSF stars sets.'
                    ' Minimized function'
                    ' is the mean of allstar\'s chi statistic, weighted by flux, calculated on '
                    ' training subset of of all stars. Results'
                    ' are stored in --dir directory if provided. List of star ids'
                    ' of the best set found, is printed to stdout (until suppressed by -no_stdout)',
        epilog=commons.version_string(),
    )
    parser.add_argument('image', default=None, nargs='?',
                        help='FITS image file (obligatory if no --demo)')
    g_dao = parser.add_argument_group('DAOPHOT/ALLSTARS related parameters',
                                      'Those parameters are used by DAOPHOT/ALLSTARS workflow. gapick optimizes PSF '
                                      'fitting errors of that workflow applied to training set stars as a function'
                                      'of \'PSF stars\' sets - stars used to derive PSF function')
    g_dao.add_argument('-c', '--all-stars-file', metavar='FILE', default=None,
                       help='all stars input file in one of daophot\'s formats (default: obtained by daophot FIND)')
    g_dao.add_argument('-D', '--daophot-opt', metavar='FILE', default=None,
                       help='daophot.opt file for daophot (default: internal)')
    g_dao.add_argument('--frames-av', metavar='n', type=int, default=1,
                       help='frames ave - parameter of daophot FIND when --all-stars-file not provided  (default: 1)')
    g_dao.add_argument('--frames-sum', metavar='n', type=int, default=1,
                       help='frames summed - parameter of daophot FIND when --all-stars-file not provided (default: 1)')
    g_dao.add_argument('-O', '--photo-opt', metavar='FILE', default=None,
                       help='photo.opt file for aperture photometry (default: none)')
    g_dao.add_argument('--photo-is', metavar='r', type=int, default=0,
                       help='PHOTOMETRY inner sky radius, overwrites photo.opt, (default: from --photo-opt or 35)')
    g_dao.add_argument('--photo-os', metavar='r', type=int, default=0,
                       help='PHOTOMETRY outher sky radius, overwrites photo.opt, (default: from --photo-opt or 50)')
    g_dao.add_argument('--photo-ap', metavar='r', type=int, default=[], nargs='+',
                       help='PHOTOMETRY apertures radius (up to 12), overwrites photo.opt, '
                            '(default: from --photo-opt or 8)')
    g_dao.add_argument('-l', '--psf-stars-file', metavar='FILE', default=None,
                       help='PSF candidates input file in one of daophot\'s formats, the result of algorithm is '
                            'a subset of those stars (default: obtained by daophot PICK)')
    g_dao.add_argument('-P', '--stars-to-pick', metavar='n', type=int, default=100,
                       help='number of stars to PICK as candidates when --stars-to-pick not provided (default: 100)')
    g_dao.add_argument('--faintest-to-pick', metavar='MAG', type=int, default=20,
                       help='faintest magnitude to PICK as candidates when --stars-to-pick not provided (default: 20)')
    g_dao.add_argument('-f', '--fine', action='store_true',
                       help='fine tuned PSF calculation (3 iter) for crowded fields, without this option no neighbours'
                            'subtraction will be performed (slower, usually not needed)')
    g_flt = parser.add_argument_group('Filtering training set stars',
                                      'Options controls which stars will be included in a training set -- whose '
                                      'residuals after the PSF fitting will be minimized by genetic algorithm')
    g_flt.add_argument('-x', '--include-psf', action='store_true',
                       help='NEW | include candidates of PSF stars in training set; by default PSF candidates are'
                            'excluded from calculation of minimized PSF error; may be included if there is low number'
                            'of stars')
    g_flt.add_argument('--max-ph-err', metavar='x', type=float, default=0.1,
                       help='threshold of aperture photometry error for training set stars; '
                            'stars for which aperture photometry (daophot PHOTO) error is greater than x '
                            'will be excluded and have no effect on minimization '
                            '(default 0.1)')
    g_flt.add_argument('--max-ph-mag', metavar='m', type=float, default=20,
                       help='threshold of aperture photometry magnitude of stars for for training set stars; '
                            '(analogous to --max-ph-err)stars for which aperture photometry (daophot PHOTO) magnitude '
                            'is greater than m (fainter than m) will be excluded form allstar run and have no effect '
                            'on quality measurement (default 20)')
    g_flt.add_argument('--clip-fit-sigmas', metavar='x', type=float, default=3.0,
                       help='NEW | threshold of errors from the preliminary PSF fitting (allstars) in sigmas above '
                            'mean error level for specific magnitude (a mean error level for the magnitude is fitted) '
                            'stars with preliminary PSF fitting error greater than mean_err(mag)+x*sigma will '
                            'be rejected from training set (default 3.0)')
    g_psf = parser.add_argument_group('Filtering candidates for PSF  stars',
                                      'Options controls which stars can be considered as PSF candidates. '
                                      'Filtering is applied to stars provided in --psf-stars-fil or found by '
                                      'daophot/PICK')
    g_psf.add_argument('--max-psf-err-mult', metavar='x', type=float, default=3.0,
                       help='threshold of PSF errors of PSF candidates - as multiplier of average PSF error; '
                            'candidates with PSF error (from daophot PSF command) greater than x*av_err will '
                            'be rejected (as well as other marked ? or * by daophot/PSF command)'
                            '(default 3.0)')
    g_psf.add_argument('--clip-psf-fit-sigmas', metavar='x', type=float, default=2.0,
                       help='NEW | threshold of errors from the preliminary PSF fitting (allstars) in sigmas above '
                            'mean error level for specific magnitude (a mean error level for the magnitude is fitted) '
                            'candidates with preliminary PSF fitting error greater than mean_err(mag)+x*sigma will '
                            'be rejected (default 2.0)')
    g_ga = parser.add_argument_group('Genetic algorithm parameters')
    g_ga.add_argument('-I', '--ga-init-prob', metavar='x', default=[0.3, 0.8], type=float, nargs='+',
                      help='what portion of candidates is used to initialize GA individuals;'
                           ' e.g. if there is 100 candidates, each of them will be '
                           ' chosen to initialize individual genome with probability x; '
                           ' in other words if x=0.3 first population in GA will contain'
                           ' individuals with around 30 stars each; '
                           ' two numbers (e.g. -I 0.3 0.8) indicates range: for each initial genome'
                           ' random number from that range will be taken (spreading number of initial stars)'
                           ' (default: 0.3 0.8)')
    g_ga.add_argument('-i', '--ga-max-iter', metavar='n', default=50, type=int,
                      help='maximum number of iterations of generic algorithm - generations (default: 50)')
    g_ga.add_argument('-n', '--ga-pop', metavar='n', default=80, type=int,
                      help='population size of GA (default: 80)')
    g_ga.add_argument('--ga-cross-prob', metavar='x', default=0.5, type=float,
                      help='crossover probability of GA (default: 0.5)')
    g_ga.add_argument('--ga-mut-prob', metavar='x', default=0.2, type=float,
                      help='mutation probability of GA - probability to became a mutant (default: 0.2)')
    g_ga.add_argument('--ga-mut-str', metavar='x', default=0.05, type=float,
                      help='mutation strength of GA - probability of every bit flip in mutant (default: 0.05)')
    g_cnt = parser.add_argument_group('gapick run control parameters')
    g_cnt.add_argument('-p', '--parallel', metavar='n', type=int, default=8,
                       help='how many parallel processes can be forked; '
                            'n=1 avoids parallelism (default: 8)')
    g_cnt.add_argument('-d', '--out-dir', metavar='output_dir', type=str, default='RESULTS',
                       help='output directory; directory will be created and result files will be stored there;'
                            ' directory should not exist or --overwrite flag should be set'
                            ' (default: RESULTS)')
    g_cnt.add_argument('-o', '--overwrite', action='store_true',
                       help='if directory specified by --out_dir parameter exists, '
                            'then ALL its content WILL BE DELETED')
    # g_cnt.add_argument('-C', '--checkpoint', metavar='file.chk', type=str, default=None,
    #                     help='restore evaluation from checkpoint; algorithm saves checkpoint.chk file every generation,'
    #                          ' which allows resuming evolution, even with another parameters')
    g_cnt.add_argument('-L', '--loglevel', metavar='level', default='info',
                       help='logging level: debug, info, warning, error, critical (default: info)')
    g_cnt.add_argument('-t', '--no-stdout', action='store_true',
                       help='suppress printing result (list of best choice of PSF stars) to stdout at finish')
    g_cnt.add_argument('-b', '--no-progress', action='store_true',
                       help='suppress showing progress bar')
    g_cnt.add_argument('-v', '--version', action='store_true',
                       help='show version and exit')
    g_cnt.add_argument('--demo', action='store_true',
                       help='use internal demo FITS file')
    __arg_parser_singlethon = parser
    return parser

def main(**kwargs):
    """Entry point for python script calls. Parameters identical to command line"""

    args = commons.bunch_kwargs(__arg_parser(), **kwargs)
    # call main routine - common form command line and python calls
    return __do(args)

def check_arguments(arg, parser=None):
    if len(arg.ga_init_prob) > 2:
        exit_error('--ga-init-prob parameter be one or two numbers', parser)
    if len(arg.ga_init_prob) == 1:
        arg.ga_init_prob = [arg.ga_init_prob[0], arg.ga_init_prob[0]]
    if max(arg.ga_init_prob) > 0.99 or min(arg.ga_init_prob) < 0.01:
        exit_error('--ga-init-prob must be in range [0.01 : 0.99]', parser)
    if arg.ga_init_prob[0] > arg.ga_init_prob[1]:
        exit_error('--ga-init-prob a b : a <= b must hold', parser)
    if not arg.demo and not arg.image:
        exit_error('FITS image parameter not provided (and no --version, --help or --demo)', parser)

def exit_error(message, parser):
    logging.error(message)
    if parser:
        parser.error()
        parser.exit(1)
    else:
        raise ValueError(message)

def info():
    """Prints commandline help message"""
    commons.info(__arg_parser())

def commandline_entry():
    # Entry point for command line
    __args = __arg_parser().parse_args()  # parse command line arguments
    if __args.version:
        print('astwro.tools ' + astwro.tools.__version__)
        exit()

    # Configure logging
    #    logging.basicConfig(format='[%(levelname)s] %(module)s: %(message)s', level=__args.loglevel.upper())
    logging.basicConfig(format='%(levelname)7s | %(message)s', level=__args.loglevel.upper())
    check_arguments(__args, __arg_parser())
    __stars = __do(__args)  # call main routine - common form command line and python calls
    if __stars is None:
        return 1
    if not __args.no_stdout:
        print('\n'.join(map(str, __stars.index)))
    return 0

if __name__ == '__main__':
    code = commandline_entry()
    exit(code)


# -*- coding: utf-8 -*-
"""
Created on Mon Dec 17 19:11:22 2018

@author: Jian
"""

import numpy as np

from . import utils

from . import BSSA14
from . import CB14
from . import CY14
from . import ASK14
from . import helper

#%%----------------------------------------------------------------------------
def GMPE(model_name, Mw, Rjb, Vs30, period, epislon=0, NGAs=None,
         rake=None, Mech=3, Ftype=None, Fnm=None, Frv=None,
         dip=None, W=None, Ztor=None, Zhypo=None, Fas=0,
         Rrup=None, Rx=None, Fhw=None, azimuth=None,
         VsFlag=0, Z25=None, Z15=None, Z10=None,
         ArbCB=0, SJ=0,
         country='California', region='CA',
         Dregion='GlobalCATW',
         CRjb=15, Ry0=None,
         D_DPP=0):

    model_name = helper.check_model_name(model_name)

    if NGAs == None:
        NGAs={'CB':{'NewCoefs':None,'terms':(1,1,1,1,1,1,1,1,1)},
              'BSSA':{'NewCoefs':None,'terms':(1,1,1)},
              'CY':{'NewCoefs':None,'terms':(1,1,1,1,1,1,1)},
              'ASK':{'NewCoefs':None,'terms':(1,1,1,1,1,1,1)}}
    else:
        if 'BSSA' not in list(NGAs.keys()):
            NGAs['BSSA'] = NGAs['BA']
        if 'ASK' not in list(NGAs.keys()):
            NGAs['ASK'] = NGAs['AS']

    dict1 = NGAs

    # check the input period
    # Note: this function is better used at a given period with a set of other parameters (not with a set of periods)
    if period > 10.0 or 0 < period < 0.01:
        raise ValueError('Positive period value should be within [0.01,10] for SA at corresponding periods')
    if period < 0 and period not in [-1,-2]:
        raise ValueError('negative period should be -1,-2 for PGA and PGV')

    if model_name == 'BSSA':
        ngaM = BSSA14.BSSA14_nga()
        kwds = {'Mech': Mech, 'Ftype': Ftype, 'Z10': Z10, 'Dregion': Dregion,
                'country': country, 'CoefTerms': dict1[model_name]}
    elif model_name == 'CB':
        ngaM = CB14.CB14_nga()
        if Fhw is None:
            Fhw = 0
        kwds = {'Ftype': Ftype, 'Rrup': Rrup, 'Ztor': Ztor, 'dip': dip,
                'Z25': Z25, 'W': W, 'Zhypo': Zhypo, 'azimuth': azimuth,
                'Fhw': Fhw, 'Z10': Z10, 'Z15': Z15, 'Arb': ArbCB, 'SJ': SJ,
                'region': region, 'CoefTerms': dict1[model_name]}
    elif model_name == 'CY':
        ngaM = CY14.CY14_nga()
        kwds = {'Ftype': Ftype, 'Rrup': Rrup, 'Rx': Rx, 'Ztor': Ztor,
                'dip': dip, 'W': W, 'Zhypo': Zhypo, 'azimuth': azimuth,
                'Fhw': Fhw, 'Z10': Z10, 'AS': Fas, 'VsFlag': VsFlag,
                'country': country, 'D_DPP': D_DPP,
                'CoefTerms': dict1[model_name]}
        # the new CY model treat PGA = SA(0.01)
        if period == -1:
            period = 0.01
    elif model_name == 'ASK':
        ngaM = ASK14.ASK14_nga()
        kwds = {'Ftype': Ftype, 'Rrup': Rrup, 'Rx': Rx, 'Ztor': Ztor,
                'dip': dip, 'W': W, 'Zhypo': Zhypo, 'azimuth': azimuth,
                'Fhw': Fhw, 'Z10': Z10, 'Fas': Fas, 'CRjb': CRjb, 'Ry0': Ry0,
                'region': region, 'country': country, 'VsFlag': VsFlag,
                'CoefTerms': dict1[model_name]}
    else:
        raise ValueError("`model_name` must be one of {'ASK', 'BSSA', 'CB', 'CY'}")

    # common interpolate for all models
    periods = np.array(ngaM.periods)
    if period in periods:
        # compute median, std directly for the existing period in the period list of the NGA model
        values = utils.mapfunc( ngaM, Mw, Rjb, Vs30, period, rake, **kwds )
        values = np.array( values )
    else:
        #print 'do the interpolation for periods that is not in the period list of the NGA model'
        ind_low = (periods <= period*1.0).nonzero()[0]
        ind_high = (periods >= period*1.0).nonzero()[0]

        period_low = max( periods[ind_low] )
        period_high = min( periods[ind_high] )
        values_low = np.array(utils.mapfunc(ngaM, Mw, Rjb, Vs30, period_low,
                                            rake, **kwds))
        values_high = np.array(utils.mapfunc(ngaM, Mw, Rjb, Vs30, period_high,
                                             rake, **kwds))
        N1,N2 = np.array(values_low).shape
        values = np.zeros( (N1,N2) )
        for icmp in range( N2 ):
            if icmp != 0:
                # stardand values are in ln (g)
                values[:,icmp] = utils.logline(np.log(period_low),
                                               np.log(period_high),
                                               values_low[:,icmp],
                                               values_high[:,icmp],
                                               np.log(period) )
            else:
                # median value is in g
                values[:,icmp] = utils.logline(np.log(period_low),
                                               np.log(period_high),
                                               np.log(values_low[:,icmp]),
                                               np.log(values_high[:,icmp]),
                                               np.log(period))
                values[:,icmp] = np.exp( values[:,icmp] )    # change the median into g unit (logline gives the result in ln(g))

    # outputs
    NGAsigmaT = values[:,1]
    NGAtau = values[:,2]
    NGAsigma = values[:,3]

    if epislon:
        NGAmedian = np.exp( np.log(values[:,0]) + epislon * NGAsigmaT )
    else:
        NGAmedian = values[:,0]

    # returned quantities are all in g, not in log(g), event for the standard deviations
    return NGAmedian[0], \
           np.exp(NGAsigmaT)[0], \
           np.exp(NGAtau)[0], \
           np.exp(NGAsigma)[0]      # all in g, include the standard deviation

#%%----------------------------------------------------------------------------
def GMPE_array(model_name, Mw, Rjb, Vs30, period_array, epislon=0, NGAs=None,
         rake=None, Mech=3, Ftype=None, Fnm=None, Frv=None,
         dip=None, W=None, Ztor=None, Zhypo=None, Fas=0,
         Rrup=None, Rx=None, Fhw=None, azimuth=None,
         VsFlag=0, Z25=None, Z15=None, Z10=None,
         ArbCB=0, SJ=0,
         country='California', region='CA',
         Dregion='GlobalCATW',
         CRjb=15, Ry0=None,
         D_DPP=0):

    import collections
    if not isinstance(period_array, collections.Iterable):
        raise TypeError('`period_array` must be a array-like object.')

    model_name = helper.check_model_name(model_name)

    results = []
    for period in period_array:
        if period < 0:
            raise ValueError('Period value must be positive')
        result_ = GMPE(model_name, Mw, Rjb, Vs30, period, epislon=epislon,
               NGAs=NGAs, rake=rake, Mech=Mech, Ftype=Ftype, Fnm=Fnm,
               Frv=Frv, dip=dip, W=W, Ztor=Ztor, Zhypo=Zhypo, Fas=Fas,
               Rrup=Rrup, Rx=Rx, Fhw=Fhw, azimuth=azimuth,
               VsFlag=VsFlag, Z25=Z25, Z15=Z15, Z10=Z10,
               ArbCB=ArbCB, SJ=SJ,country=country, region=region,
               Dregion=Dregion, CRjb=CRjb, Ry0=Ry0, D_DPP=D_DPP)

        results.append(result_)

    results_T = list(zip(*results))

    return tuple([np.array(_) for _ in results_T])


import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
from utils import utils_sr
import torch
from argparse import ArgumentParser
from utils.utils_restoration import rgb2y, psnr, array2tensor, tensor2array
import sys
from matplotlib.ticker import MaxNLocator

class PnP_restoration():

    def __init__(self, hparams):

        self.hparams = hparams
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.initialize_cuda_denoiser()

    def initialize_cuda_denoiser(self):
        '''
        Initialize the denoiser model with the given pretrained ckpt
        '''
        sys.path.append('../GS_denoising/')
        from lightning_denoiser import GradMatch
        parser2 = ArgumentParser(prog='utils_restoration.py')
        parser2 = GradMatch.add_model_specific_args(parser2)
        parser2 = GradMatch.add_optim_specific_args(parser2)
        hparams = parser2.parse_known_args()[0]
        if 'nb_4' in self.hparams.pretrained_checkpoint :
            hparams.DRUNET_nb = 4
        hparams.grad_matching = self.hparams.grad_matching
        hparams.act_mode = 's'
        self.denoiser_model = GradMatch(hparams)
        checkpoint = torch.load(self.hparams.pretrained_checkpoint, map_location=self.device)
        self.denoiser_model.load_state_dict(checkpoint['state_dict'])
        self.denoiser_model.eval()
        for i, v in self.denoiser_model.named_parameters():
            v.requires_grad = False
        self.denoiser_model = self.denoiser_model.to(self.device)
        if self.hparams.precision == 'double' :
            if self.denoiser_model is not None:
                self.denoiser_model.double()

    def initialize_prox(self, img, degradation):
        '''
        calculus for future prox computatations
        :param img: degraded image
        :param degradation: 2D blur kernel for deblurring and SR, mask for inpainting
        '''
        if self.hparams.degradation_mode == 'deblurring' :
            self.k = degradation
            self.k_tensor = array2tensor(np.expand_dims(self.k, 2)).double().to(self.device)
            self.FB, self.FBC, self.F2B, self.FBFy = utils_sr.pre_calculate_prox(img, self.k_tensor, 1)
        elif self.hparams.degradation_mode == 'SR':
            self.k = degradation
            self.k_tensor = array2tensor(np.expand_dims(self.k, 2)).double().to(self.device)
            self.FB, self.FBC, self.F2B, self.FBFy = utils_sr.pre_calculate_prox(img, self.k_tensor, 2)
        elif self.hparams.degradation_mode == 'inpainting':
            self.M = array2tensor(degradation).double().to(self.device)
            self.My = self.M*img
        else:
            print('degradation mode not treated')

    def calculate_prox(self, img):
        '''
        Calculation of the proximal mapping of the data term f
        :param img: input for the prox
        :return: prox_f(img)
        '''
        if self.hparams.degradation_mode == 'deblurring':
            rho = torch.tensor([1/self.hparams.lamb]).double().repeat(1, 1, 1, 1).to(self.device)
            px = utils_sr.prox_solution(img.double(), self.FB, self.FBC, self.F2B, self.FBFy, rho, 1)
        elif self.hparams.degradation_mode == 'SR':
            rho = torch.tensor([1 /self.hparams.lamb]).double().repeat(1, 1, 1, 1).to(self.device)
            px = utils_sr.prox_solution(img.double(), self.FB, self.FBC, self.F2B, self.FBFy, rho, self.hparams.sf)
        elif self.hparams.degradation_mode == 'inpainting':
            if self.hparams.noise_level_img > 1e-2:
                px = (self.hparams.lamb*self.My + img)/(self.hparams.lamb*self.M+1)
            else :
                px = self.My + (1-self.M)*img
        else:
            print('degradation mode not treated')
        return px

    def calculate_grad(self, img):
        '''
        Calculation of the gradient of the data term f
        :param img: input for the prox
        :return: \nabla_f(img)
        '''
        if self.hparams.degradation_mode == 'deblurring' :
            grad = utils_sr.grad_solution(img.double(), self.FB, self.FBC, self.FBFy, 1)
        if self.hparams.degradation_mode == 'SR' :
            grad = utils_sr.grad_solution(img.double(), self.FB, self.FBC, self.FBFy, self.hparams.sf)
        return grad

    def calculate_regul(self,y,x,g):
        '''
        Calculation of the regularization (1/tau)*phi_sigma(y)
        :param y: Point where to evaluate
        :param x: D^{-1}(y)
        :param g: Precomputed regularization function value at x
        :return: regul(y)
        '''
        regul = (1 / self.hparams.lamb) * (g - (1 / 2) * torch.norm(x - y, p=2) ** 2)
        return regul

    def calulate_data_term(self,y,img):
        '''
        Calculation of the data term value f(y)
        :param y: Point where to evaluate F
        :param img: Degraded image
        :return: f(y)
        '''
        if self.hparams.degradation_mode == 'deblurring':
            deg_y = utils_sr.imfilter(y.double(), self.k_tensor[0].double().flip(1).flip(2).expand(3, -1, -1, -1))
            f = 0.5 * torch.norm(img - deg_y, p=2) ** 2
        elif self.hparams.degradation_mode == 'SR':
            deg_y = utils_sr.imfilter(y.double(), self.k_tensor[0].double().flip(1).flip(2).expand(3, -1, -1, -1))
            deg_y = deg_y[..., 0::self.hparams.sf, 0::self.hparams.sf]
            f = 0.5 * torch.norm(img - deg_y, p=2) ** 2
        elif self.hparams.degradation_mode == 'inpainting':
            deg_y = self.M * y.double()
            f = 0.5 * torch.norm(img - deg_y, p=2) ** 2
        else:
            print('degradation not implemented')
        return f

    def calculate_F(self, y, x, g, img):
        '''
        Calculation of the objective function value f(y) + (1/tau)*phi_sigma(y)
        :param y: Point where to evaluate F
        :param x: D^{-1}(y)
        :param g: Precomputed regularization function value at x
        :param img: Degraded image
        :return: F(y)
        '''
        regul = self.calculate_regul(y,x,g)
        if self.hparams.no_data_term:
            F = regul
        else:
            f = self.calulate_data_term(y,img)
            F = f + regul
        return F.item()

    def calculate_lyapunov_DRS(self,y,z,x,g,img):
        '''
            Calculation of the Lyapunov function value Psi(x)
            :param x: Point where to evaluate F
            :param y,z: DRS iterations initialized at x
            :param g: Precomputed regularization function value at x
            :param img: Degraded image
            :return: Psi(x)
        '''
        regul = self.calculate_regul(y,x,g)
        f = self.calulate_data_term(z, img)
        Psi = regul + f + (1 / self.hparams.lamb) * (torch.sum(torch.mul(y-x,y-z)) + (1/2) * torch.norm(y - z, p=2) ** 2)
        return Psi

    def restore(self, img, init_im, clean_img, degradation,extract_results=False):
        '''
        Compute GS-PnP restoration algorithm
        :param img: Degraded image
        :param clean_img: ground-truth clean image
        :param degradation: 2D blur kernel for deblurring and SR, mask for inpainting
        :param extract_results: Extract information for subsequent image or curve saving
        '''

        if extract_results:
            y_list, z_list, x_list, Dg_list, psnr_tab, g_list, Dx_list, F_list, Psi_list = [], [], [], [], [], [], [], [], []

        i = 0 # iteration counter

        img_tensor = array2tensor(init_im).to(self.device) # for GPU computations (if GPU available)
        self.initialize_prox(img_tensor, degradation)  # prox calculus that can be done outside of the loop

        # Initialization of the algorithm
        if self.hparams.degradation_mode == 'SR':
            x0 = cv2.resize(init_im, (img.shape[1] * self.hparams.sf, img.shape[0] * self.hparams.sf),interpolation=cv2.INTER_CUBIC)
            x0 = utils_sr.shift_pixel(x0, self.hparams.sf)
            x0 = array2tensor(x0).to(self.device)
        else:
            x0 = array2tensor(init_im).to(self.device)

        if extract_results:  # extract np images and PSNR values
            out_x = tensor2array(x0.cpu())
            current_x_psnr = psnr(clean_img, out_x)
            if self.hparams.print_each_step:
                print('current x PSNR : ', current_x_psnr)
            psnr_tab.append(current_x_psnr)
            x_list.append(out_x)

        x = x0

        if self.hparams.use_hard_constraint:
            x = torch.clamp(x, 0, 1)

        # Initialize Lyapunov
        diff_Psi = 1
        Psi_old = 1
        Psi = Psi_old

        while i < self.hparams.maxitr and abs(diff_Psi)/Psi_old > self.hparams.relative_diff_Psi_min:

            if self.hparams.inpainting_init :
                if i < self.hparams.n_init:
                    self.sigma_denoiser = 50
                else :
                    self.sigma_denoiser = self.hparams.sigma_denoiser
            else :
                self.sigma_denoiser = self.hparams.sigma_denoiser

            x_old = x
            Psi_old = Psi

            if self.hparams.PnP_algo == 'PGD':
                # Gradient step
                gradx = self.calculate_grad(x_old)
                z = x_old - self.hparams.lamb*gradx
                # Denoising step
                torch.set_grad_enabled(True)
                Dg, N = self.denoiser_model.calculate_grad(z, self.hparams.sigma_denoiser / 255.)
                torch.set_grad_enabled(False)
                Dg = Dg.detach()
                N = N.detach()
                g = 0.5 * (torch.norm(z.double() - N.double(), p=2) ** 2)
                Dz = z - Dg
                Dx = Dz
                x = (1 - self.hparams.alpha) * z + self.hparams.alpha*Dz
                y = x
                # Hard constraint
                if self.hparams.use_hard_constraint:
                    x = torch.clamp(x,0,1)
                # Calculate Objective
                F = self.calculate_F(x, z, g, img_tensor)

            elif self.hparams.PnP_algo == 'DRS':
                # Denoising step
                torch.set_grad_enabled(True)
                Dg, N = self.denoiser_model.calculate_grad(x_old, self.hparams.sigma_denoiser / 255.)
                torch.set_grad_enabled(False)
                Dg = Dg.detach()
                N = N.detach()
                g = 0.5 * (torch.norm(x_old.double() - N.double(), p=2) ** 2)
                Dx = x_old - Dg
                y = (1 - self.hparams.alpha)*x_old + self.hparams.alpha*Dx
                # Hard constraint
                if self.hparams.use_hard_constraint:
                    y = torch.clamp(y,0,1)
                # Proximal step
                z = self.calculate_prox(2*y-x_old)
                # Calculate Lyapunov
                Psi = self.calculate_lyapunov_DRS(y,z,x,g,img_tensor)
                diff_Psi = Psi-Psi_old
                # Calculate Objective
                F = self.calculate_F(y, x, g, img_tensor)
                # Final step
                x = x_old + (z-y)

            elif self.hparams.PnP_algo == 'DRSdiff':

                # Proximal step
                y = self.calculate_prox(x_old)
                y2 = 2*y-x_old

                # Denoising step
                torch.set_grad_enabled(True)
                Dg, N = self.denoiser_model.calculate_grad(y2, self.hparams.sigma_denoiser / 255.)
                torch.set_grad_enabled(False)
                Dg = Dg.detach()
                N = N.detach()
                g = 0.5 * (torch.norm(y2.double() - N.double(), p=2) ** 2)
                Dx = y2 - Dg
                z = (1 - self.hparams.alpha) * y2 + self.hparams.alpha * Dx
                # Hard constraint
                if self.hparams.use_hard_constraint:
                    z = torch.clamp(z, 0, 1)

                # Calculate Lyapunov
                Psi = self.calculate_lyapunov_DRS(y,z,x,g,img_tensor)
                diff_Psi = Psi-Psi_old
                # Calculate Objective
                F = self.calculate_F(y, x, g, img_tensor)
                # Final step
                x = x_old + (z-y)

            else :
                print('algo not implemented')

            # Logging
            if extract_results:
                out_y = tensor2array(y.cpu())
                out_z = tensor2array(z.cpu())
                out_x = tensor2array(x.cpu())
                current_y_psnr = psnr(clean_img, out_y)
                current_z_psnr = psnr(clean_img, out_z)
                current_x_psnr = psnr(clean_img, out_x)
                if self.hparams.print_each_step:
                    print('iteration : ', i)
                    print('current y PSNR : ', current_y_psnr)
                    print('current z PSNR : ', current_z_psnr)
                    print('current x PSNR : ', current_x_psnr)
                y_list.append(out_y)
                x_list.append(out_x)
                z_list.append(out_z)
                Dx_list.append(tensor2array(Dx.cpu()))
                Dg_list.append(torch.norm(Dg).cpu().item())
                g_list.append(g.cpu().item())
                psnr_tab.append(current_x_psnr)
                F_list.append(F)
                Psi_list.append(Psi)

            # next iteration
            i += 1

        output_img = tensor2array(y.cpu())
        output_psnr = psnr(clean_img, output_img)
        output_psnrY = psnr(rgb2y(clean_img), rgb2y(output_img))

        if extract_results:
            return output_img, output_psnr, output_psnrY, x_list, np.array(z_list), np.array(y_list), np.array(Dg_list), np.array(psnr_tab), np.array(Dx_list), np.array(g_list), np.array(F_list), np.array(Psi_list)
        else:
            return output_img, output_psnr, output_psnrY

    def initialize_curves(self):

        self.conv = []
        self.PSNR = []
        self.g = []
        self.Dg = []
        self.F = []
        self.Psi = []
        self.lip_algo = []
        self.lip_D = []
        self.lip_Dg = []

    def update_curves(self, x_list, Dx_list, psnr_tab, Dg_list, g_list, F_list, Psi_list):

        self.F.append(F_list)
        self.Psi.append(Psi_list)
        self.g.append(g_list)
        self.Dg.append(Dg_list)
        self.PSNR.append(psnr_tab)
        self.conv.append(np.array([(np.linalg.norm(x_list[k + 1] - x_list[k]) ** 2) for k in range(len(x_list) - 1)]) / np.sum(np.abs(x_list[0]) ** 2))
        self.lip_algo.append(np.sqrt(np.array([np.sum(np.abs(x_list[k + 1] - x_list[k]) ** 2) for k in range(1, len(x_list) - 1)]) / np.array([np.sum(np.abs(x_list[k] - x_list[k - 1]) ** 2) for k in range(1, len(x_list[:-1]))])))
        self.lip_D.append(np.sqrt(np.array([np.sum(np.abs(Dx_list[i + 2] - Dx_list[i+1]) ** 2) for i in range(len(Dx_list) - 2)]) / np.array([np.sum(np.abs(x_list[i+1] - x_list[i]) ** 2) for i in range(len(Dx_list) - 2)])))
        self.lip_Dg.append(np.sqrt(np.array([np.sum(np.abs(Dg_list[i + 2] - Dg_list[i+1]) ** 2) for i in range(len(Dx_list) - 2)]) / np.array([np.sum(np.abs(x_list[i+1] - x_list[i]) ** 2) for i in range(len(Dg_list) - 2)])))

    def save_curves(self, save_path):

        import matplotlib
        matplotlib.rcParams.update({'font.size': 10})
        matplotlib.rcParams['lines.linewidth'] = 2
        matplotlib.style.use('seaborn-darkgrid')

        plt.figure(0)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.g)):
            plt.plot(self.g[i], markevery=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.legend()
        plt.savefig(os.path.join(save_path, 'g.png'),bbox_inches="tight")

        plt.figure(1)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.PSNR)):
            plt.plot(self.PSNR[i], markevery=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.legend()
        plt.savefig(os.path.join(save_path, 'PSNR.png'),bbox_inches="tight")

        plt.figure(2)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.F)):
            plt.plot(self.Psi[i], markevery=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.legend()
        plt.savefig(os.path.join(save_path, 'Liapunov.png'), bbox_inches="tight")

        plt.figure(22)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.F)):
            plt.plot(self.F[i], markevery=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.legend()
        plt.savefig(os.path.join(save_path, 'F.png'), bbox_inches="tight")

        plt.figure(4)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.Dg)):
            Ds_norm = [np.linalg.norm(np.array(self.Dg[i][j])) for j in range(len(self.Dg[i]))]
            plt.plot(Ds_norm, linewidth=1.5, markevery=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.legend()
        plt.savefig(os.path.join(save_path, 'Dg.png'), bbox_inches="tight")

        plt.figure(5)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.conv)):
            plt.plot(self.conv[i], '-o', markevery=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.legend()
        plt.savefig(os.path.join(save_path, 'conv_log.png'), bbox_inches="tight")

        self.conv2 = [[np.min(self.conv[i][:k]) for k in range(1, len(self.conv[i]))] for i in range(len(self.conv))]
        conv_rate = [self.conv2[i][0]*np.array([(1/k) for k in range(1,len(self.conv2[i]))]) for i in range(len(self.conv2))]
        plt.figure(6)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.conv)):
            plt.plot(self.conv2[i], '-', markevery=10)
        plt.plot(conv_rate[i], '--', color='red', label=r'$\mathcal{O}(\frac{1}{K})$')
        plt.semilogy()
        plt.legend()
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.savefig(os.path.join(save_path, 'conv_log2.png'), bbox_inches="tight")

        self.conv_sum = [[np.sum(self.conv[i][:k]) for k in range(1, len(self.conv[i]))] for i in range(len(self.conv))]

        plt.figure(7)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.conv_sum)):
            plt.plot(self.conv_sum[i], markevery=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.legend()
        plt.savefig(os.path.join(save_path, 'conv_log_sum.png'), bbox_inches="tight")

        plt.figure(8)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.lip_algo)):
            plt.plot(self.lip_algo[i], '-o', markersize=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.savefig(os.path.join(save_path, 'lip_algo.png'))

        plt.figure(9)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.lip_D)):
            plt.plot(self.lip_D[i], '-o', markersize=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.savefig(os.path.join(save_path, 'lip_D.png'))

        plt.figure(10)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.lip_Dg)):
            plt.plot(self.lip_Dg[i], '-o', markersize=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.savefig(os.path.join(save_path, 'lip_Dg.png'))

    def add_specific_args(parent_parser):
        parser = ArgumentParser(parents=[parent_parser], add_help=False)
        parser.add_argument('--dataset_path', type=str, default='../datasets')
        parser.add_argument('--pretrained_checkpoint', type=str,default='../GS_denoising/ckpts/Prox_DRUNet.ckpt')
        parser.add_argument('--PnP_algo', type=str, default='PGD')
        parser.add_argument('--dataset_name', type=str, default='CBSD68')
        parser.add_argument('--sigma_denoiser', type=float)
        parser.add_argument('--sigma_k_denoiser', type=float)
        parser.add_argument('--noise_level_img', type=float, default=2.55)
        parser.add_argument('--maxitr', type=int, default=1000)
        parser.add_argument('--alpha', type=float, default=1)
        parser.add_argument('--lamb', type=float)
        parser.add_argument('--n_images', type=int, default=68)
        parser.add_argument('--relative_diff_Psi_min', type=float, default=1e-8)
        parser.add_argument('--inpainting_init', dest='inpainting_init', action='store_true')
        parser.set_defaults(inpainting_init=False)
        parser.add_argument('--precision', type=str, default='simple')
        parser.add_argument('--n_init', type=int, default=10)
        parser.add_argument('--patch_size', type=int, default=256)
        parser.add_argument('--extract_curves', dest='extract_curves', action='store_true')
        parser.set_defaults(extract_curves=False)
        parser.add_argument('--extract_images', dest='extract_images', action='store_true')
        parser.set_defaults(extract_images=False)
        parser.add_argument('--print_each_step', dest='print_each_step', action='store_true')
        parser.set_defaults(print_each_step=False)
        parser.add_argument('--no_grad_matching', dest='grad_matching', action='store_false')
        parser.set_defaults(grad_matching=True)
        parser.add_argument('--no_data_term', dest='no_data_term', action='store_true')
        parser.set_defaults(no_data_term=False)
        parser.add_argument('--use_hard_constraint', dest='use_hard_constraint', action='store_true')
        parser.set_defaults(no_data_term=False)
        return parser


# test protocols for setting up relative binding free energy calculations.

import argparse
import os
import numpy as np
from rdkit import Chem
from rdkit.Chem import AllChem

from timemachine.lib import potentials, custom_ops
from timemachine.lib import LangevinIntegrator

from ff.handlers import openmm_deserializer
from ff.handlers.deserialize import deserialize_handlers

from fe import pdb_writer
from fe import rbfe
from md import Recipe
from md import builders

from multiprocessing import Pool

def get_romol_conf(mol):
    conformer = mol.GetConformer(0)
    guest_conf = np.array(conformer.GetPositions(), dtype=np.float64)
    guest_conf = guest_conf/10 # from angstroms to nm
    return np.array(guest_conf, dtype=np.float64)

def run(args):

    lamb, intg, bound_potentials, masses, x0, box, gpu_idx, stage = args
    # print("running on", gpu_idx)
    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_idx)
    u_impls = []
    for bp in bound_potentials:
        u_impls.append(bp.bound_impl(precision=np.float32))

    # important that we reseed here.
    intg.seed = np.random.randint(np.iinfo(np.int32).max)
    intg_impl = intg.impl()

    v0 = np.zeros_like(x0)

    ctxt = custom_ops.Context(
        x0,
        v0,
        box,
        intg_impl,
        u_impls
    )

    # secondary minimization needed only for stage 1
    if stage == 1:
        for l in np.linspace(0.35, lamb, 500):
            ctxt.step(l)

    # equilibration
    for step in range(20000):
    # for step in range(1000):
        ctxt.step(lamb)

    # print(ctxt.get_x_t())

    du_dl_obs = custom_ops.AvgPartialUPartialLambda(u_impls, 10)
    ctxt.add_observable(du_dl_obs)

    # add observable for <du/dl>
    for step in range(50000):
    # for step in range(5000):
        ctxt.step(lamb)

    print(lamb, du_dl_obs.avg_du_dl())

    assert np.any(np.abs(ctxt.get_x_t()) > 100) == False
    assert np.any(np.isnan(ctxt.get_x_t())) == False
    assert np.any(np.isinf(ctxt.get_x_t())) == False

    return du_dl_obs.avg_du_dl()

def minimize(args):

    bound_potentials, masses, x0, box = args

    u_impls = []
    for bp in bound_potentials:
        u_impls.append(bp.bound_impl(precision=np.float32))

    seed = np.random.randint(np.iinfo(np.int32).max)

    intg = LangevinIntegrator(
        300.0,
        1.5e-3,
        1.0,
        masses,
        seed
    ).impl()

    v0 = np.zeros_like(x0)

    ctxt = custom_ops.Context(
        x0,
        v0,
        box,
        intg,
        u_impls
    )

    steps = 500

    lambda_schedule = np.linspace(0.35, 0.0, 500)
    for lamb in lambda_schedule:
        ctxt.step(lamb)

    return ctxt.get_x_t()

def minimize_setup(r_host, r_ligand):
    r_combined = r_host.combine(r_ligand)

    print(len(r_combined.masses))

    # assert 0
    host_atom_idxs = np.arange(len(r_host.masses))
    ligand_atom_idxs = np.arange(len(r_ligand.masses)) + len(r_host.masses)
    rbfe.set_nonbonded_lambda_idxs(r_combined, ligand_atom_idxs, 0, 1)

    # for bp in r_combined.bound_potentials:
        # if isinstance(bp, potentials.Nonbonded):
            # print(len(bp.get_lambda_offset_idxs()))
            # print(len(bp.get_lambda_plane_idxs()))
            # assert 0

    return r_combined

def main(args, stage):

    # benzene = Chem.AddHs(Chem.MolFromSmiles("c1ccccc1")) # a
    # phenol = Chem.AddHs(Chem.MolFromSmiles("Oc1ccccc1")) # b
                                            #01234567890
    benzene = Chem.AddHs(Chem.MolFromSmiles("C1=CC=C2C=CC=CC2=C1")) # a
    phenol =  Chem.AddHs(Chem.MolFromSmiles("C1=CC=C2C=CC=CC2=C1")) # b

    AllChem.EmbedMolecule(benzene)
    AllChem.EmbedMolecule(phenol)

    ff_handlers = deserialize_handlers(open('ff/params/smirnoff_1_1_0_ccc.py').read())
    r_benzene = Recipe.from_rdkit(benzene, ff_handlers)
    r_phenol = Recipe.from_rdkit(phenol, ff_handlers)

    r_combined = r_benzene.combine(r_phenol)
    core_pairs = np.array([
        [0,0],
        [1,1],
        [2,2],
        [3,3],
        [4,4],
        [5,5],
        [6,6],
        [7,7],
        [8,8],
        [9,9],
        # [10,10]
    ], dtype=np.int32)
    core_pairs[:, 1] += benzene.GetNumAtoms()

    a_idxs = np.arange(benzene.GetNumAtoms())
    b_idxs = np.arange(phenol.GetNumAtoms()) + benzene.GetNumAtoms()

    core_k = 20.0

    if stage == 0:
        centroid_k = 200.0
        rbfe.stage_0(r_combined, b_idxs, core_pairs, centroid_k, core_k)
        # lambda_schedule = np.linspace(0.0, 1.0, 2)
        # lambda_schedule = np.array([0.0, 0.0, 0.0, 0.0, 0.0])
        lambda_schedule = np.array([0.0, 0.0, 0.0, 0.0, 0.0])
    elif stage == 1:
        rbfe.stage_1(r_combined, a_idxs, b_idxs, core_pairs, core_k)
        lambda_schedule = np.linspace(0.0, 1.2, 60)
    else:
        assert 0

    system, host_coords, box, topology = builders.build_water_system(4.0)

    r_host = Recipe.from_openmm(system)
    r_final = r_host.combine(r_combined)

    # minimize coordinates of host + ligand A
    ha_coords = np.concatenate([
        host_coords,
        get_romol_conf(benzene)
    ])

    pool = Pool(args.num_gpus)

    # we need to run this in a subprocess since the cuda runtime
    # must not be initialized in the master thread due to lack of
    # fork safety
    r_minimize = minimize_setup(r_host, r_benzene)
    ha_coords = pool.map(minimize, [(r_minimize.bound_potentials, r_minimize.masses, ha_coords, box)], chunksize=1)
    # this is a list
    ha_coords = ha_coords[0]
    pool.close()

    pool = Pool(args.num_gpus)

    x0 = np.concatenate([
        ha_coords,
        get_romol_conf(phenol)
    ])

    masses = np.concatenate([r_host.masses, r_benzene.masses, r_phenol.masses])

    seed = np.random.randint(np.iinfo(np.int32).max)

    intg = LangevinIntegrator(
        300.0,
        1.5e-3,
        1.0,
        masses,
        seed
    )

    # production run at various values of lambda
    for epoch in range(10):
        avg_du_dls = []

        run_args = []
        for lamb_idx, lamb in enumerate(lambda_schedule):
            run_args.append((
                lamb,
                intg,
                r_final.bound_potentials,
                r_final.masses,
                x0,
                box,
                lamb_idx % args.num_gpus,
                stage))

        avg_du_dls = pool.map(run, run_args, chunksize=1)

        print("stage", stage, "epoch", epoch, "dG", np.trapz(avg_du_dls, lambda_schedule))

if __name__ == "__main__":

    parser = argparse.ArgumentParser(
        description="RBFE testing",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument(
        "--num_gpus",
        type=int,
        help="number of gpus"
    )

    args = parser.parse_args()

    main(args, 0)


import copy
import os
import numpy as np
import torch
import torch.nn.functional as F
from common.buffers import OfflineBuffer
from utils.train_tools import soft_target_update, evaluate
from utils import log_tools
from common.networks import MLPSquashedReparamGaussianPolicy, CVAE, MLPQsaNet

class BEAR_Agent:
    """
    Implementation of Bootstrapping Error Accumulation Reduction (BEAR)
    https://arxiv.org/abs/1906.00949
    BEAR's MMD Loss's weight alpha_prime is tuned automatically by default.

    Actor Loss: alpha_prime * MMD Loss + -minQ(s,a)
    Critic Loss: Like BCQ
    Alpha_prime Loss: -(alpha_prime * (MMD Loss - threshold))
    """
    def __init__(self,
                 env,
                 data_buffer: OfflineBuffer,
                 policy_net: MLPSquashedReparamGaussianPolicy,  # actor
                 q_net1: MLPQsaNet,  # critic
                 q_net2: MLPQsaNet,
                 cvae_net: CVAE,
                 policy_lr=1e-4,
                 qf_lr=3e-4,
                 cvae_lr=3e-4,
                 gamma=0.99,
                 tau=0.05,

                 # BEAR
                 lmbda=0.75,  # used for double clipped double q-learning
                 mmd_sigma=20.0,  # the sigma used in mmd kernel
                 kernel_type='gaussian',  # the type of mmd kernel(gaussian or laplacian)
                 lagrange_thresh=0.05,  # the hyper-parameter used in automatic tuning alpha in cql loss
                 n_action_samples=100,  # the number of action samples to compute the best action when choose action
                 n_target_samples=10,  # the number of action samples to compute BCQ-like target value
                 n_mmd_action_samples=4,  # the number of action samples to compute MMD.
                 warmup_step=40000,  # do support matching with a warm start before policy(actor) train

                 max_train_step=1000000,
                 log_interval=1000,
                 eval_freq=5000,
                 train_id="bear_hopper-medium-v2_test",
                 resume=False,  # if True, train from last checkpoint
                 device='cpu',
                 ):

        self.env = env
        self.data_buffer = data_buffer

        self.device = torch.device(device)

        # the network and optimizers
        self.policy_net = policy_net.to(self.device)
        self.q_net1 = q_net1.to(self.device)
        self.q_net2 = q_net2.to(self.device)
        self.target_q_net1 = copy.deepcopy(self.q_net1).to(self.device)
        self.target_q_net2 = copy.deepcopy(self.q_net2).to(self.device)
        self.cvae_net = cvae_net.to(self.device)
        self.policy_optimizer = torch.optim.Adam(self.policy_net.parameters(), lr=policy_lr)
        self.q_optimizer1 = torch.optim.Adam(self.q_net1.parameters(), lr=qf_lr)
        self.q_optimizer2 = torch.optim.Adam(self.q_net2.parameters(), lr=qf_lr)
        self.cvae_optimizer = torch.optim.Adam(self.cvae_net.parameters(), lr=cvae_lr)

        self.gamma = gamma
        self.tau = tau

        self.max_train_step = max_train_step
        self.eval_freq = eval_freq
        self.train_step = 0

        self.resume = resume  # whether load checkpoint start train from last time

        # BEAR
        self.lmbda = lmbda
        self.mmd_sigma = mmd_sigma
        self.kernel_type = kernel_type
        self.lagrange_thresh = lagrange_thresh
        self.n_action_samples = n_action_samples
        self.n_target_samples = n_target_samples
        self.n_mmd_action_samples = n_mmd_action_samples
        self.warmup_step = warmup_step

        # mmd loss's temperature
        self.log_alpha_prime = torch.zeros(1, requires_grad=True, device=self.device)
        self.alpha_prime_optimizer = torch.optim.Adam([self.log_alpha_prime], lr=1e-3)

        # log dir and interval
        self.log_interval = log_interval
        self.result_dir = os.path.join(log_tools.ROOT_DIR, "run/results", train_id)
        log_tools.make_dir(self.result_dir)
        self.checkpoint_path = os.path.join(self.result_dir, "checkpoint.pth")
        self.tensorboard_writer = log_tools.TensorboardLogger(self.result_dir)

    def choose_action(self, obs, eval=False):
        with torch.no_grad():
            obs = torch.FloatTensor(obs).reshape(1, -1).repeat(self.n_action_samples, 1).to(self.device)
            action, _, _ = self.policy_net(obs)
            q1 = self.q_net1(obs, action)
            ind = q1.argmax(dim=0)
        return action[ind].cpu().numpy().flatten()

    def mmd_loss_laplacian(self, samples1, samples2, sigma=0.2):
        """MMD constraint with Laplacian kernel for support matching"""
        # sigma is set to 10.0 for hopper, cheetah and 20 for walker/ant
        diff_x_x = samples1.unsqueeze(2) - samples1.unsqueeze(1)  # B x N x N x d
        diff_x_x = torch.mean((-(diff_x_x.abs()).sum(-1)/(2.0 * sigma)).exp(), dim=(1, 2))

        diff_x_y = samples1.unsqueeze(2) - samples2.unsqueeze(1)
        diff_x_y = torch.mean((-(diff_x_y.abs()).sum(-1)/(2.0 * sigma)).exp(), dim=(1, 2))

        diff_y_y = samples2.unsqueeze(2) - samples2.unsqueeze(1)  # B x N x N x d
        diff_y_y = torch.mean((-(diff_y_y.abs()).sum(-1)/(2.0 * sigma)).exp(), dim=(1, 2))

        overall_loss = (diff_x_x + diff_y_y - 2.0 * diff_x_y + 1e-6).sqrt()
        return overall_loss

    def mmd_loss_gaussian(self, samples1, samples2, sigma=0.2):
        """MMD constraint with Gaussian Kernel support matching"""
        # sigma is set to 10.0 for hopper, cheetah and 20 for walker/ant
        diff_x_x = samples1.unsqueeze(2) - samples1.unsqueeze(1)  # B x N x N x d
        diff_x_x = torch.mean((-(diff_x_x.pow(2)).sum(-1)/(2.0 * sigma)).exp(), dim=(1, 2))

        diff_x_y = samples1.unsqueeze(2) - samples2.unsqueeze(1)
        diff_x_y = torch.mean((-(diff_x_y.pow(2)).sum(-1)/(2.0 * sigma)).exp(), dim=(1, 2))

        diff_y_y = samples2.unsqueeze(2) - samples2.unsqueeze(1)  # B x N x N x d
        diff_y_y = torch.mean((-(diff_y_y.pow(2)).sum(-1)/(2.0 * sigma)).exp(), dim=(1, 2))

        overall_loss = (diff_x_x + diff_y_y - 2.0 * diff_x_y + 1e-6).sqrt()
        return overall_loss

    def train(self):
        # Sample
        batch = self.data_buffer.sample()
        obs = batch["obs"].to(self.device)
        acts = batch["acts"].to(self.device)
        rews = batch["rews"].to(self.device)
        next_obs = batch["next_obs"].to(self.device)
        done = batch["done"].to(self.device)

        """
        Train the Behaviour cloning policy to be able to take more than 1 sample for MMD.
        Conditional VAE is used as Behaviour cloning policy in BEAR.
        """
        recon_action, mu, log_std = self.cvae_net(obs, acts)
        cvae_loss = self.cvae_net.loss_function(recon_action, acts, mu, log_std)

        self.cvae_optimizer.zero_grad()
        cvae_loss.backward()
        self.cvae_optimizer.step()

        """
        Critic Training
        """
        with torch.no_grad():
            # generate 10 actions for every next_obs(Same as BCQ)
            next_obs = torch.repeat_interleave(next_obs, repeats=self.n_target_samples, dim=0).to(self.device)
            # compute target Q value of generated action
            target_q1 = self.target_q_net1(next_obs, self.policy_net(next_obs)[0])
            target_q2 = self.target_q_net2(next_obs, self.policy_net(next_obs)[0])
            # soft clipped double q-learning
            target_q = self.lmbda * torch.min(target_q1, target_q2) + (1. - self.lmbda) * torch.max(target_q1, target_q2)
            # take max over each action sampled from the generation and perturbation model
            target_q = target_q.reshape(obs.shape[0], self.n_target_samples, 1).max(1)[0].squeeze(1)
            target_q = rews + self.gamma * (1. - done) * target_q

        # compute current Q
        current_q1 = self.q_net1(obs, acts).squeeze(1)
        current_q2 = self.q_net2(obs, acts).squeeze(1)
        # compute critic loss
        critic_loss1 = F.mse_loss(current_q1, target_q)
        critic_loss2 = F.mse_loss(current_q2, target_q)

        self.q_optimizer1.zero_grad()
        critic_loss1.backward()
        self.q_optimizer1.step()

        self.q_optimizer2.zero_grad()
        critic_loss2.backward()
        self.q_optimizer2.step()

        # MMD Loss
        # sample actions from dataset and current policy(B x N x D)
        raw_sampled_actions = self.cvae_net.decode_multiple_without_squash(obs, decode_num=self.n_mmd_action_samples,
                                                                           z_device=self.device)
        raw_actor_actions = self.policy_net.sample_multiple_without_squash(obs, sample_num=self.n_mmd_action_samples)
        if self.kernel_type == 'gaussian':
            mmd_loss = self.mmd_loss_gaussian(raw_sampled_actions, raw_actor_actions, sigma=self.mmd_sigma)
        else:
            mmd_loss = self.mmd_loss_laplacian(raw_sampled_actions, raw_actor_actions, sigma=self.mmd_sigma)

        """
        Alpha prime training(lagrangian parameter update for MMD loss weight)
        """
        alpha_prime_loss = -(self.log_alpha_prime.exp() * (mmd_loss - self.lagrange_thresh)).mean()
        self.alpha_prime_optimizer.zero_grad()
        alpha_prime_loss.backward(retain_graph=True)
        self.alpha_prime_optimizer.step()

        self.log_alpha_prime.data.clamp_(min=-5.0, max=10.0)  # clip for stability

        """
        Actor Training
        Actor Loss = alpha_prime * MMD Loss + -minQ(s,a)
        """
        a, log_prob, _ = self.policy_net(obs)
        min_q = torch.min(self.q_net1(obs, a), self.q_net2(obs, a)).squeeze(1)
        # policy_loss = (self.alpha * log_prob - min_q).mean()  # SAC Type
        policy_loss = - (min_q.mean())

        # BEAR Actor Loss
        actor_loss = (self.log_alpha_prime.exp() * mmd_loss).mean()
        if self.train_step > self.warmup_step:
            actor_loss = policy_loss + actor_loss
        self.policy_optimizer.zero_grad()
        actor_loss.backward()  # the mmd_loss will backward again in alpha_prime_loss.
        self.policy_optimizer.step()

        soft_target_update(self.q_net1, self.target_q_net1, tau=self.tau)
        soft_target_update(self.q_net2, self.target_q_net2, tau=self.tau)

        self.train_step += 1

        return critic_loss1.cpu().item(), critic_loss2.cpu().item(), policy_loss.cpu().item(), alpha_prime_loss.cpu().item()

    def learn(self):
        if self.resume:
            self.load_agent_checkpoint()
        else:
            # delete tensorboard log file
            log_tools.del_all_files_in_dir(self.result_dir)

        while self.train_step < (int(self.max_train_step)):
            # train
            q_loss1, q_loss2, policy_loss, alpha_prime_loss = self.train()

            if self.train_step % self.eval_freq == 0:
                avg_reward, avg_length = evaluate(agent=self, episode_num=5)
                self.tensorboard_writer.log_eval_data({"eval_episode_length": avg_length,
                                                       "eval_episode_reward": avg_reward}, self.train_step)

            if self.train_step % self.log_interval == 0:
                self.store_agent_checkpoint()
                self.tensorboard_writer.log_train_data({"q_loss_1": q_loss1,
                                                        "q_loss_2": q_loss2,
                                                        "policy_loss": policy_loss,
                                                        "alpha_prime_loss": alpha_prime_loss}, self.train_step)

    def store_agent_checkpoint(self):
        checkpoint = {
            "q_net1": self.q_net1.state_dict(),
            "q_net2": self.q_net2.state_dict(),
            "policy_net": self.policy_net.state_dict(),
            "q_optimizer1": self.q_optimizer1.state_dict(),
            "q_optimizer2": self.q_optimizer2.state_dict(),
            "policy_optimizer": self.policy_optimizer.state_dict(),
            "log_alpha_prime": self.log_alpha_prime,
            "alpha_prime_optimizer": self.alpha_prime_optimizer.state_dict(),
            "train_step": self.train_step,
        }

        torch.save(checkpoint, self.checkpoint_path)

    def load_agent_checkpoint(self):
        checkpoint = torch.load(self.checkpoint_path, map_location=self.device)  # can load gpu's data on cpu machine
        self.q_net1.load_state_dict(checkpoint["q_net1"])
        self.q_net2.load_state_dict(checkpoint["q_net2"])
        self.policy_net.load_state_dict(checkpoint["policy_net"])
        self.q_optimizer1.load_state_dict(checkpoint["q_optimizer1"])
        self.q_optimizer2.load_state_dict(checkpoint["q_optimizer2"])
        self.policy_optimizer.load_state_dict(checkpoint["policy_optimizer"])
        self.log_alpha_prime = checkpoint["log_alpha_prime"]
        self.alpha_prime_optimizer.load_state_dict(checkpoint["alpha_prime_optimizer"])
        self.train_step = checkpoint["train_step"]

        print("load checkpoint from \"" + self.checkpoint_path +
              "\" at " + str(self.train_step) + " time step")