#
# Copyright (C) 2019 Kevin Thornton <krthornt@uci.edu>
#
# This file is part of fwdpy11.
#
# fwdpy11 is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# fwdpy11 is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with fwdpy11.  If not, see <http://www.gnu.org/licenses/>.
#
"""
Local adaptation of a quantitative trait to differing optima.
"""

import fwdpy11
import pandas as pd
import math
import numpy as np
import sys
import argparse
from collections import namedtuple

# Simulations with tree sequence recording need
# to know the max position in a genome.  Here,
# we use a length of 1.0. Thus, all mutation
# and recombination events will be uniform
# random variables on the continuous interval
# [0, GENOME_LENGTH).
GENOME_LENGTH = 1.0

# When recording quant-genetic statistics during a simulation,
# we will use this type. Named tuples are extremely efficient,
# and they are easily converted into Pandas DataFrame objects,
# which is very convenient for analysis and output.
Datum = namedtuple('Data', ['generation', 'deme', 'gbar', 'vg', 'wbar'])


def make_parser():
    """
    Create a command-line interface to the script.
    """
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    required = parser.add_argument_group("Required arguments")
    required.add_argument("--popsize", "-N", type=int,
                          help="Diploid population size")
    required.add_argument("--mu", "-m", type=float,
                          help="Mutation rate (per gamete, per generation)")
    required.add_argument("--sigma", "-s", type=float,
                          help="Standard deviation of Gaussian"
                          "distribution of mutational effects")

    optional = parser.add_argument_group("Optional arguments")
    optional.add_argument("--rho", type=float, default=1000.0,
                          help="Scaled recombination rate, rho=4Nr")
    optional.add_argument("--VS", type=float, default=10.0,
                          help="Inverse strength of stabilizing selection")
    optional.add_argument("--opt", type=float, default=1.0,
                          help="Value of new phenotypic optimum")
    optional.add_argument("--migrates", type=float, nargs=2, default=None,
                          help="Migration rates from 0 to 1 and 1 to 0, respectively.")
    optional.add_argument("--time", type=float, default=0.1,
                          help="Amount of time to simulate past"
                          "optimum shift, in units of N")
    optional.add_argument("--plotfile", type=str, default=None,
                          help="File name for plot")
    optional.add_argument("--seed", type=int, default=42,
                          help="Random number seed.")

    return parser


def validate_arguments(args):
    """
    Validate input arguments.
    Note: this is likely incomplete.
    """
    if args.popsize is None:
        raise ValueError("popsize cannot be None")
    if args.mu < 0:
        raise ValueError("mu must be non-negative")
    if args.mu is None:
        raise ValueError("mu cannot be None")
    if args.mu < 0 or math.isfinite(args.mu) is False:
        raise ValueError("Mutation rate must be non-negative and finite")
    if args.sigma is None:
        raise ValueError("sigma cannot be none")
    if args.sigma < 0 or math.isfinite(args.sigma) is False:
        raise ValueError("Std. dev. of distribution of effect sizes"
                         "must be non-negative and finite")

    if args.migrates is not None:
        for m in args.migrates:
            if m < 0 or m > 1:
                raise ValueError("migration rates must be 0 <= m <= 1")


def make_migmatrix(migrates):
    if migrates is None:
        return None
    mm = np.zeros(4).reshape(2, 2)
    mm[0, 1] = migrates[1]
    mm[1, 0] = migrates[0]
    rs = np.sum(mm, axis=1)
    np.fill_diagonal(mm, 1.0 - rs)
    return fwdpy11.MigrationMatrix(mm)


class Recorder(object):
    """
    fwdpy11 allows you to define objects that record data
    from populations during simulation.  Such objects must
    be callable, and the easiest way to do things is to
    create a class with a __call__ function.
    """

    def __init__(self, start):
        self.data = []
        self.start = start

    def __call__(self, pop, recorder):
        if pop.generation >= self.start:
            # Record mean trait value each generation.
            md = np.array(pop.diploid_metadata, copy=False)
            demes = np.unique(md['deme'])
            for d in demes:
                w = np.where(md['deme'] == d)[0]
                gbar = md['g'][w].mean()
                vg = md['g'][w].var()
                wbar = md['w'][w].mean()
                self.data.append(Datum(pop.generation, d,
                                       gbar, vg, wbar))


def plot_output(data, filename):
    import matplotlib.pyplot as plt
    import matplotlib.gridspec as gridspec
    fig = plt.figure(figsize=(9, 3))
    gs = gridspec.GridSpec(ncols=3, nrows=1, figure=fig)
    ax_gbar = fig.add_subplot(gs[0, 0])
    ax_vg = fig.add_subplot(gs[0, 1])
    ax_wbar = fig.add_subplot(gs[0, 2])

    df = pd.DataFrame(data, columns=Datum._fields)
    g = df.groupby(['deme'])
    for n, gi in g:
        ax_gbar.plot(gi['generation'], gi['gbar'], label="Deme {}".format(n))
        ax_vg.plot(gi['generation'], gi['vg'], label="Deme {}".format(n))
        ax_wbar.plot(gi['generation'], gi['wbar'], label="Deme {}".format(n))

    for ax in [ax_gbar, ax_vg, ax_wbar]:
        ax.set_xlabel("Generation")

    ax_gbar.set_ylabel(r'$\bar{g}$')
    ax_vg.set_ylabel(r'$V(G)$')
    ax_wbar.set_ylabel(r'$\bar{w}$')
    ax_gbar.legend()
    plt.tight_layout()
    plt.savefig(filename)


def runsim(args):
    """
    Run the simulation.
    """
    pop = fwdpy11.DiploidPopulation(2*args.popsize, GENOME_LENGTH)

    np.random.seed(args.seed)
    rng = fwdpy11.GSLrng(args.seed)

    GSSmo0 = fwdpy11.GSSmo(
        [(0, 0, args.VS), (10*args.popsize, args.opt, args.VS)])
    GSSmo1 = fwdpy11.GSSmo(
        [(0, 0, args.VS), (10*args.popsize, -1.0*args.opt, args.VS)])

    mm = make_migmatrix(args.migrates)
    dd = fwdpy11.DiscreteDemography(
        mass_migrations=[fwdpy11.move_individuals(0, 0, 1, 0.5)],
        migmatrix=mm)

    p = {'nregions': [],  # No neutral mutations -- add them later!
         'gvalue': [fwdpy11.Additive(2.0, GSSmo0),
                    fwdpy11.Additive(2.0, GSSmo1)],
         'sregions': [fwdpy11.GaussianS(0, GENOME_LENGTH, 1, args.sigma)],
         'recregions': [fwdpy11.Region(0, GENOME_LENGTH, 1)],
         'rates': (0.0, args.mu, args.rho/float(4*args.popsize)),
         # Keep mutations at frequency 1 in the pop if they affect fitness.
         'prune_selected': False,
         'demography': dd,
         'simlen': 10*args.popsize + int(args.popsize*args.time)
         }
    params = fwdpy11.ModelParams(**p)

    r = Recorder(10*args.popsize)
    fwdpy11.evolvets(rng, pop, params, 100, r, suppress_table_indexing=True)
    if args.plotfile is not None:
        plot_output(r.data, args.plotfile)


if __name__ == "__main__":
    parser = make_parser()
    args = parser.parse_args(sys.argv[1:])
    validate_arguments(args)
    runsim(args)
