#!/usr/bin/env python

"""skynet.py: Image segnemtation algoritm for Solar Disk: SkyNet"""

__author__ = "Chakraborty, S."
__copyright__ = ""
__credits__ = []
__license__ = "MIT"
__version__ = "1.0."
__maintainer__ = "Chakraborty, S."
__email__ = "shibaji7@vt.edu"
__status__ = "Research"

import matplotlib.pyplot as plt
import matplotlib as mpl
from dateutil import parser as prs
import pandas as pd
import datetime as dt
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.autograd import Variable
import cv2
import sys
import numpy as np
import torch.nn.init
import random
import os
import json

_params_ = {
    "_desc": "PyTorch Unsupervised Segmentation",
    "nChannel": "Number of channels",
    "maxIter": "Number of maximum iterations",
    "minLabels": "Minimum number of labels",
    "lr": "Learning rate",
    "nConv": "Number of convolutional layers",
    "nConv": "Number of convolutional layers",
    "stepsize_sim": "Step size for similarity loss",
    "stepsize_con": "Step size for continuity loss",
}

import astropy.units as u
from sunpy.net import Fido, attrs
import sunpy.map
from aiapy.calibrate import register, update_pointing, normalize_exposure

import matplotlib.pyplot as plt
import cv2
import os

class RegisterAIA(object):

    def __init__(self, date, wavelength=193, resolution=1024, vmin=10):
        self.wavelength = wavelength
        self.date = date
        self.resolution = resolution
        self.vmin = vmin
        self.folder = "data/SDO-Database/{:4d}.{:02d}.{:02d}/{:04d}/{:03d}/".format(self.date.year, self.date.month, self.date.day,
                                                                             self.resolution, self.wavelength)
        self.fname = "{:4d}_{:02d}_{:02d}_{:02d}{:02d}{:02d}.png".format(self.date.year, self.date.month,
                                                                         self.date.day, self.date.hour,
                                                                         self.date.minute, self.date.second)
        if not os.path.exists(self.folder): os.system("mkdir -p " + self.folder)
        if not os.path.exists(self.folder + self.fname):
            self.normalized()
            self.to_png()
        return

    def normalized(self):
        q = Fido.search(
            attrs.Time(self.date.strftime("%Y-%m-%dT%H:%M:%S"), (self.date + dt.timedelta(seconds=11)).strftime("%Y-%m-%dT%H:%M:%S")),
            attrs.Instrument("AIA"),
            attrs.Wavelength(wavemin=self.wavelength*u.angstrom, wavemax=self.wavelength*u.angstrom),
        )
        self.m = sunpy.map.Map(Fido.fetch(q[0,0]))
        m_updated_pointing = update_pointing(self.m)
        m_registered = register(m_updated_pointing)
        self.m_normalized = normalize_exposure(m_registered)
        return

    def to_png(self):
        norm, m = self.m_normalized, self.m
        fig, ax = plt.subplots(nrows=1,ncols=1,dpi=100,figsize=(2048/100, 2048/100))
        norm.plot(annotate=False, axes=ax, vmin=self.vmin)
        ax.set_xticks([])
        ax.set_yticks([])
        fig.savefig("tmp.png",bbox_inches="tight")
        im = cv2.imread("tmp.png")
        im = im[10:-10,10:-10]
        im = cv2.resize(im, (self.resolution, self.resolution))
        os.remove("tmp.png")
        cv2.imwrite(self.folder + self.fname, im)
        return

# CNN model
class SkyNet(nn.Module):

    def __init__(self, input_dim, params):
        super(SkyNet, self).__init__()
        self.params = params
        self.conv1 = nn.Conv2d(input_dim, self.params.nChannel, kernel_size=3, stride=1, padding=1 )
        self.bn1 = nn.BatchNorm2d(self.params.nChannel)
        self.conv2 = nn.ModuleList()
        self.bn2 = nn.ModuleList()
        for i in range(args.nConv-1):
            self.conv2.append( nn.Conv2d(self.params.nChannel, self.params.nChannel, kernel_size=3, stride=1, padding=1 ) )
            self.bn2.append( nn.BatchNorm2d(self.params.nChannel) )
        self.conv3 = nn.Conv2d(self.params.nChannel, self.params.nChannel, kernel_size=1, stride=1, padding=0 )
        self.bn3 = nn.BatchNorm2d(self.params.nChannel)
        return

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu( x )
        x = self.bn1(x)
        for i in range(self.params.nConv-1):
            x = self.conv2[i](x)
            x = F.relu( x )
            x = self.bn2[i](x)
        x = self.conv3(x)
        x = self.bn3(x)
        return x

class Loader(object):

    def __init__(self, fname, folder, date, params, save=True, cfg_file = "data/config/{:3d}.json"):
        self.fname = fname
        self.folder = folder
        self.setup(cfg_file, params)
        self.date = date
        self.params = params
        self.resolution = params.resolution
        self.wavelength = params.wavelength
        self.save = save
        self.extn = "." + fname.split(".")[-1]
        self.use_cuda = torch.cuda.is_available()
        self.detect_hough_circles()
        self.im = cv2.imread(self.folder + self.fname)
        self.data = torch.from_numpy( np.array([self.im.transpose( (2, 0, 1) ).astype("float32")/255.]) )
        if self.use_cuda: self.data = self.data.cuda()
        self.data = Variable(self.data)
        return

    def rescale(self, img, to):
        """ Resacle the images """
        dsize = (to, to)
        img = cv2.resize(img, dsize)
        return img

    def setup(self, cfg_file, params):
        """ Load parameters """
        _dict_ = {}
        for k in vars(params).keys():
            _dict_[k] = vars(args)[k]
        cfg_file = cfg_file.format(_dict_["wavelength"])
        with open(cfg_file, "r") as fp:
            dic = json.load(fp)
            for p in dic.keys():
                if not hasattr(self, p): setattr(self, p, dic[p])
        for p in _dict_.keys():
            setattr(self, p, _dict_[p])
        self._dict_ = _dict_
        mult = int(self.resolution/512) + 0.5
        self.prims = np.array(self.prims) * mult
        self.delta = int(self.resolution/128)

        self.images = {}
        self.images["src"] = cv2.imread(cv2.samples.findFile(self.folder+self.fname), cv2.IMREAD_COLOR)
        self.img_params = {"resolution": self.images["src"].shape[0]}
        self.images["src"] = self.rescale(self.images["src"], self.resolution)
        self.images["org"] = np.copy(cv2.cvtColor(self.images["src"], cv2.COLOR_BGR2RGB))
        self.images["prob_masked_gray_image"], self.images["prob_masked_image"] =\
        np.zeros_like(self.images["src"]), np.copy(self.images["src"])
        self.images["gray"] = cv2.cvtColor(self.images["src"], cv2.COLOR_BGR2GRAY)
        self.images["contours"] = np.zeros_like(self.images["gray"])
        self.images["contour_maps"] = np.zeros_like(self.images["gray"])
        self.images["bin_map"] = np.zeros_like(self.images["gray"])
        self.images["hc_mask"] = np.zeros_like(self.images["gray"])
        self.images["prob_masked_image_ol"] = np.copy(self.images["src"])
        return

    def detect_hough_circles(self):
        """ Detecting the Hough Circles """
        gray = np.copy(self.images["gray"])
        rows = self.resolution
        circles = cv2.HoughCircles(gray, cv2.HOUGH_GRADIENT, 1, rows / 8,
                                   param1=self.hc_param1, param2=self.hc_param2,
                                   minRadius=int(rows/3), maxRadius=int(rows/2))
        if circles is not None:
            circles = np.uint16(np.around(circles))
            radi = int(rows/2)
            for i in circles[0, :]:
                if radi - self.alpha <= i[0] <= radi + self.alpha and radi - self.alpha <= i[1] <= radi + self.alpha:
                    self.center = (i[0], i[1])
                    self.radius = i[2]
                    self.create_hc_mask()
        return

    def create_hc_mask(self):
        # HC mask create
        mask = cv2.circle(self.images["hc_mask"], self.center, self.radius, 255, -1)
        self.images["hc_mask"] = mask
        self.NzC = np.count_nonzero(mask)
        return

    def load_model(self):
        # Load & train model
        self.model = SkyNet( self.data.size(1), self.params)
        if self.use_cuda: self.model.cuda()
        self.model.train()
        self.loss_fn = torch.nn.CrossEntropyLoss()
        # scribble loss definition
        loss_fn_scr = torch.nn.CrossEntropyLoss()
        # continuity loss definition
        self.loss_hpy = torch.nn.L1Loss(size_average = True)
        self.loss_hpz = torch.nn.L1Loss(size_average = True)

        self.HPy_target = torch.zeros(self.im.shape[0]-1, self.im.shape[1], self.params.nChannel)
        self.HPz_target = torch.zeros(self.im.shape[0], self.im.shape[1]-1, self.params.nChannel)
        if self.use_cuda: self.HPy_target, self.HPz_target = self.HPy_target.cuda(), self.HPz_target.cuda()
        self.optimizer = optim.SGD(self.model.parameters(), lr=args.lr, momentum=0.9)
        self.label_colours = np.random.randint(255, size=(100,3))
        return self

    def run_forwarding(self):
        if not os.path.exists(self.folder + self.fname.replace(self.extn, "_seg" + self.extn)):
            for batch_idx in range(self.params.maxIter):
                self.optimizer.zero_grad()
                self.output = self.model( self.data )[ 0 ]
                self.output = self.output.permute( 1, 2, 0 ).contiguous().view( -1, self.params.nChannel )

                self.outputHP = self.output.reshape( (self.im.shape[0], self.im.shape[1], self.params.nChannel) )
                self.HPy = self.outputHP[1:, :, :] - self.outputHP[0:-1, :, :]
                self.HPz = self.outputHP[:, 1:, :] - self.outputHP[:, 0:-1, :]
                self.lhpy = self.loss_hpy(self.HPy, self.HPy_target)
                self.lhpz = self.loss_hpz(self.HPz, self.HPz_target)

                ignore, target = torch.max( self.output, 1 )
                im_target = target.data.cpu().numpy()
                nLabels = len(np.unique(im_target))
                self.loss = self.params.stepsize_sim * self.loss_fn(self.output, target) + self.params.stepsize_con * (self.lhpy + self.lhpz)
                self.loss.backward()
                self.optimizer.step()
                print (batch_idx, "/", self.params.maxIter, "|", " label num :", nLabels, " | loss : %.2f"%self.loss.item())
                if nLabels <= self.params.minLabels or self.params.los > self.loss:
                    print (" >> nLabels:", nLabels, " reached minLabels:", self.params.minLabels, " with loss: %.2f."%self.loss.item())
                    break
        return self

    def save_outputs(self):
        if not os.path.exists(self.folder + self.fname.replace(self.extn, "_seg" + self.extn)):
            output = self.model(self.data)[0]
            output = output.permute(1, 2, 0).contiguous().view(-1, self.params.nChannel)
            ignore, target = torch.max(output, 1)
            self.im_target = target.data.cpu().numpy()
            self.unique_targets = np.array([self.label_colours[c % 100] for c in np.unique(self.im_target)])
            self.im_target_rgb = np.array([self.label_colours[c % 100] for c in self.im_target])
            self.im_target_rgb = self.im_target_rgb.reshape(self.im.shape).astype(np.uint8)
            cv2.imwrite(self.folder + self.fname.replace(self.extn, "_seg" + self.extn), self.im_target_rgb)
        return self

    def check_intensity(self, mask):
        # Check the intensity bound on the output
        ints = False
        gray = np.copy(self.images["gray"])
        gray = cv2.bitwise_and(gray, gray, mask=mask)
        if np.quantile(gray.ravel(), 0.05) >= 0 and np.quantile(gray.ravel(), 0.95) <= 100: ints = True
        return ints

    def estimate_CHB(self):
        if not os.path.exists(self.folder + self.fname.replace(self.extn, "_msk" + self.extn)):
            maskList = []
            for targ in self.unique_targets:
                mask = cv2.inRange(self.im_target_rgb, targ-1, targ+1)
                mask = cv2.bitwise_and(mask, mask, mask=self.images["hc_mask"])
                if np.count_nonzero(mask) < 0.1*self.NzC and self.check_intensity(mask):
                    maskList.append(mask)
            self.totalmask = np.array(maskList).sum(axis=0)
            cv2.imwrite(self.folder + self.fname.replace(self.extn, "_msk" + self.extn), self.totalmask)
        else:
            fname = self.folder + self.fname.replace(self.extn, "_msk" + self.extn)
            self.totalmask = cv2.cvtColor(cv2.imread(fname, cv2.IMREAD_COLOR), cv2.COLOR_BGR2GRAY)
        return self

    def rescale_prob(self, prob, l=0., u=0.5):
        prob = prob*(u-l) + l
        return prob

    def calculate_probabilistic_boundaries(self):
        self.contours, self.hierarchy = cv2.findContours((255-self.totalmask).astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        for ix in range(len(self.contours)):
            if ix > 0: self.draw_contour(self.contours[int(ix)], self.images["prob_masked_image"])
        cv2.imwrite(self.folder + self.fname.replace(self.extn, "_stg0.0_out" + self.extn), self.images["prob_masked_image"])
        gray, mask = np.copy(self.images["gray"]), np.zeros_like(self.images["gray"])
        contours = self.contours[1:]
        cv2.drawContours(mask, contours, -1, 255, -1); mask = cv2.bitwise_and(gray, gray, mask=mask)
        maskplot = np.copy(mask)
        cv2.imwrite(self.folder + self.fname.replace(self.extn, "_stg0.1_out" + self.extn), mask)
        thds = [16, 24, 32, 48, 64, 96, 128, 255]
        linked_list = []
        for td in thds:
            masked = np.zeros_like(mask)
            masked[mask <= td] = 255
            masked = cv2.bitwise_and(masked, masked, mask=mask)
            blurs = cv2.blur((255 - masked).astype(np.uint8), (3, 3))
            contours, hierarchy = cv2.findContours(blurs, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
            for ix in range(len(contours)):
                if ix > 0 and hierarchy[0,ix,3] == 0:
                    cimg = np.zeros_like(gray)
                    cv2.drawContours(cimg, contours, int(ix), color=255, thickness=-1)
                    pts = np.where(cimg == 255)
                    intensity = self.intensity_threshold-gray[pts[0], pts[1]].ravel().astype(int)
                    prob = 1-np.mean(1./(1.+np.exp(intensity)))
                    colw = np.array([255,255,255])
                    #print(td, "->", prob)
                    linked_list.append({"prob":prob, "color":255*prob, "contour":contours[int(ix)]})
                    cv2.drawContours(maskplot, contours, int(ix), color=255*prob, thickness=1)
                    self.draw_contour(contours[int(ix)], self.images["prob_masked_image_ol"], col=tuple(colw*prob), line_thick=1)
            cv2.imwrite(self.folder + self.fname.replace(self.extn, "_stg0.2.%03d_out" + self.extn)%td, masked)
        cv2.imwrite(self.folder + self.fname.replace(self.extn, "_stg1.0_out" + self.extn), self.images["prob_masked_image_ol"])
        maskplot = cv2.cvtColor(maskplot, cv2.COLOR_GRAY2RGB)
        maskplot = cv2.applyColorMap(maskplot, cv2.COLORMAP_JET)
        cv2.imwrite(self.folder + self.fname.replace(self.extn, "_stg1.1_out" + self.extn), maskplot)
        self.plot_image_dev(self.images["prob_masked_image_ol"], linked_list)
        self.plot_image_dev(maskplot, linked_list, cmap=mpl.cm.jet, ext="_cb_jet")
        return self

    def plot_image_dev(self, image, objs, shape=(15,14), cmap=mpl.cm.gray, ext="_cb_gray"):
        df = pd.DataFrame.from_records(objs)[["prob","color"]]
        fig = plt.figure(figsize=(4, 4), dpi=180)
        ax = plt.subplot2grid(shape, (0,0), rowspan=shape[0]-2, colspan=shape[0]-2)
        ax.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        ax.set_xticks([])
        ax.set_yticks([])
        norm = mpl.colors.Normalize(vmin=0, vmax=1)
        ax = plt.subplot2grid(shape, (2,shape[0]-2), rowspan=9, colspan=1)
        cb1 = mpl.colorbar.ColorbarBase(ax, cmap=cmap,
                                        norm=norm,
                                        orientation="vertical")
        cb1.set_label("Pr(CH)")
        fig.savefig(self.folder + self.fname.replace(self.extn, ext + self.extn), bbox_inches="tight")
        return

    def draw_contour(self, contour, img, gray=False, col=None, line_thick=None):
        if line_thick is None: line_thick = self.draw_param["contur"]["thick"]
        if col is None: col=tuple(self.draw_param["contur"]["color"])
        for _x in range(contour.shape[0]-1):
            if gray: cv2.line(img, (contour[_x,0,0], contour[_x,0,1]), (contour[_x+1,0,0], contour[_x+1,0,1]), (255,255,255), line_thick)
            else: cv2.line(img, (contour[_x,0,0], contour[_x,0,1]), (contour[_x+1,0,0], contour[_x+1,0,1]), col, line_thick)
        if gray: cv2.line(img, (contour[0,0,0], contour[0,0,1]), (contour[-1,0,0], contour[-1,0,1]), (255,255,255), line_thick)
        else: cv2.line(img, (contour[0,0,0], contour[0,0,1]), (contour[-1,0,0], contour[-1,0,1]),  col, line_thick)
        return

    def save_files_outputs(self):
        return self

def run_skynet(args, save=True):
    aia = RegisterAIA(args.date, args.wavelength, args.resolution, vmin=10)
    load = Loader(aia.fname, aia.folder, args.date, args, save)
    load.load_model().run_forwarding().save_outputs().estimate_CHB()
    load.calculate_probabilistic_boundaries()
    if save: load.save_files_outputs()
    return

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="PyTorch Unsupervised Segmentation")
    parser.add_argument("-dn", "--date", default=dt.datetime(2018,5,30,12), help="Date [2018,5,30,12]", type=prs.parse)
    parser.add_argument("-r", "--resolution", default=1024, help="Resolution of the files [1024]", type=int)
    parser.add_argument("-w", "--wavelength", default=193, help="Wavelength of the files [193]", type=int)
    parser.add_argument("--nChannel", metavar="N", default=50, type=int, help="number of channels")
    parser.add_argument("--maxIter", metavar="T", default=100, type=int, help="number of maximum iterations")
    parser.add_argument("--minLabels", metavar="minL", default=3, type=int, help="minimum number of labels")
    parser.add_argument("--lr", metavar="LR", default=0.3, type=float, help="learning rate")
    parser.add_argument("--nConv", metavar="M", default=2, type=int, help="number of convolutional layers")
    parser.add_argument("--stepsize_sim", metavar="SIM", default=1, type=float, help="step size for similarity loss", required=False)
    parser.add_argument("--stepsize_con", metavar="CON", default=1, type=float, help="step size for continuity loss")
    parser.add_argument("--los", metavar="LOS", default=.1, type=float, help="Final loss value")
    args = parser.parse_args()
    print("\n Parameter list for SkyNet simulation ")
    for k in vars(args).keys():
        print("     " + k + "->" + str(vars(args)[k]))
    run_skynet(args)


"""
gc.head() offers functions for calling specific simulation scenarios and
producing plots and processed data files from the simulation results. Contains
functionality better suited for smaller as well as larger simulations as well
as evaluations of specific parameter changes or vaccination scenarios.
"""

from __future__ import division
from builtins import range
import math
import numpy as np
import pandas as pd
import scipy.stats
import time
import seaborn
from collections import Counter
from importlib import *

import cf
import gc_memo
import gc_plots

reload(cf)
reload(gc_memo)
reload(gc_plots)

from gc_memo import main
from gc_plots import *

def small_scale(store_export='dictionary'):
    """ Performs a single simulation of the given system, creates
    the following plots:
            - population dynamics overview (free naive cells, free memory cells
            and GC populations over time)
            - for each GC, a clonal composition plot together with its memory
            output in a separate panel
            - for each GC, the evolution of its largest clone's affinities over
            mutation count (this plot contains aritificial noise to increase
            visibility!).

    If store_export is set 'datafile', the simulation data is stored in a
    hdf5 file for future purposes, for 'dictionary' the data is passed
    internally and discarded after the run.

    Recommended only for small simulation sizes with up to ~5 GCs and ~5k
    cells, as otherwise things get crowded and plots get large.
    """
    # get runID from current system time
    runID = int(time.time())
    # run simulation and get filepath or dict
    simdata = main(runID, store_export=store_export, evalperday=12)
    # import required information for small scale plots
    l_times, l_fn, l_fm, l_GCs, LFcurve, Agcurve, evaltimes, freePan, GCPans, \
        ms_times, ms_vals, ms_fams, ms_muts, mut_list, \
        E_list = import_file(simdata)
    # plot population behaviour
    population_plot(l_times, l_fn, l_fm, l_GCs, runID)
    # plot GC contents and memory output for every GC
    for i in range(len(l_GCs)):
        GC_dynamics_plot(GCPans[i], ms_times[i], ms_fams[i], ms_vals[i],
                         ms_muts[i], runID, i)

    return(simdata)

def TUCHMI_sampling(store_export='datafile', d_export=True, subsample=12):
    """ Performs a single simulation of a given size using the TUCHMI vaccination
    protocol, samples memory from the simulated pool and creates several plots
    summarising the information. User settings regarding the protocol are
    overwritten.

    Plots produced include:
        - mean SHMs and clonal expansion (fraction of cells sampled from
        clones that appeared more than once within the sample) in samples of
        size subsample
        - scatter plot of affinity over mutational status in polyclonal samples
        at TUCHMI time points I, II and III
        - scatter plot of affinity over mutational status at a clonal level,
        cells sampled from three TUCHMI time points merged into single plots
        (but sampling time point encoded in colouring)

    If store_export is set 'datafile', the simulation data is stored in a
    hdf5 file for future purposes, for 'dictionary' the data is passed
    internally and lost after the run.

    If d_export is set True, textfiles containing the sampled data (used for
    plotting) are exported for each plot individually.

    Subsample gives the number of cells to be sampled at each timepoint in
    oder to calculate entropy and unique fraction.

    Can be used for all simulation sizes, but is especially useful for larger
    simulations (e.g. >=50 GCs, 50k cells).
    """
    # give protocol
    cf.endtime = 126*12
    cf.tinf = [0*12, 28*12, 56*12]
    cf.dose = [1, 1, 1]

    # get runID from current system time
    runID = int(time.time())

    # run simulation and get filepath or dict
    simdata = main(runID, store_export=store_export, evalperday=1)

    # import required information
    l_times, l_fn, l_fm, l_GCs, LFcurve, Agcurve, evaltimes, freePan, GCPans, \
        ms_times, ms_vals, ms_fams, ms_muts, mut_list, E_list = \
        import_file(simdata)

    # for affinity-mutation scatter plot, downsample for visibility
    pick_tp = 35*12
    samplefrac = 100./len(freePan.sel(timepoint = pick_tp).dropna("dim_0"))

    # get list of lists to catch values at every timepoint
    tList = list(freePan["timepoint"].values)
    TT = len(tList)
    SHM_means = [[] for t in range(TT)]
    Entropies = [[] for t in range(TT)]
    clusterfracs = [[] for t in range(TT)]

    """ Cell pool affinity over time """
    # get mean affinity at all time points
    Elist = []
    for tp in range(len(tList)):
        C = freePan.sel(timepoint = tList[tp]).loc[dict(dim_1="affinity")].dropna("dim_0").values.mean()
        Elist.append(C)
    # pass energies to plot function
    pool_affinity_plot(tList, Elist)

    """ Mean SHM and clonal expansion within sample of size subsample """
    # sample 100 times to calculate standard deviations
    for nn in range(100):
        for tp in range(TT):
            ttp = 12*tp
            # cellnumber to be sampled is either subsample or, if less cells
            # are available (more of a hypothetic case really), all cells
            freePan_no_na = freePan.sel(timepoint = ttp).dropna("dim_0")
            cellnum = min(subsample, len(freePan_no_na))
            if cellnum > 0:
                cell_id = np.random.choice(len(freePan_no_na), cellnum)
                cells = freePan_no_na[cell_id, :]
                c_muts = list(cells.loc[dict(dim_1="mutations")].values)
                SHM_means[tp].append(np.nanmean(c_muts))
                # evaluate entropies and clusterfractions
                CC = Counter(list(cells.loc[dict(dim_1="family")].values))
                Entropies[tp].append(scipy.stats.entropy(list(CC.values()), base=2)
                                     / math.log(cellnum, 2))
                # count again to find how many clones have one member only,
                # calculate clusterfrac from this
                sizedist = list(CC.values())
                C2 = Counter(sizedist)
                uniquefrac = float(C2[1])/cellnum
                clusterfracs[tp].append(1-uniquefrac)
            else:
                SHM_means[tp].append(np.nan)
                Entropies[tp].append(np.nan)
                uniquefrac[tp].append(np.nan)
                clusterfracs[tp].append(np.nan)

    # pass information to plotting function
    MSHM = np.nanmean(SHM_means, axis=1)
    SSHM = np.nanstd(SHM_means, axis=1)
    MEntropies = np.nanmean(Entropies, axis=1)
    SEntropies = np.nanstd(Entropies, axis=1)
    Mclusterfracs = np.nanmean(clusterfracs, axis=1)
    Sclusterfracs = np.nanstd(clusterfracs, axis=1)

    sample_statistics_plot(subsample, tList, MSHM, SSHM,
                           MEntropies, SEntropies, Mclusterfracs,
                           Sclusterfracs)

    """ Plot of affinity/mutations on tps I, II and III """
    # sample cells 7 days post each infection, record SHM, KD and origin
    # (memory versus naive first activated ancestor)
    timecourse = [7*12, 35*12, 63*12]
    SHM_list = [[] for t in timecourse]
    KD_list = [[] for t in timecourse]
    orglist = [[] for t in timecourse]

    for d in range(len(timecourse)):
        tp = timecourse[d]
        freePan_no_na = freePan.sel(timepoint = tp).dropna("dim_0")
        cellnum = int(np.round(len(freePan_no_na)*samplefrac))
        if cellnum > 0:
            cell_id = np.random.choice(len(freePan_no_na), cellnum)
            cells = freePan_no_na[cell_id, :]
            kdl = list(cells.loc[dict(dim_1="affinity")].values)
            # transform norm E to KD
            kdll = np.exp(cf.y0+np.array(kdl)*cf.m)
            KD_list[d] = list(kdll)
            # get mutation counts, correct them and origin
            SHM_list[d] = list(cells.loc[dict(dim_1="mutations")].values)
            orglist[d] = list(cells.loc[dict(dim_1="origin")].values)

    # pass information to plot function
    sample_scatter_plot(KD_list, SHM_list, orglist)

    """ Affinity/mutation plots for individual clusters """
    # samples from the memory pool at the given timepoints, split information
    # into clusters and plot SHM/KD scatter plots for some of these clusters.

    # lists to collect SHM, KD values, families and timepoints for all panels
    SHM_list = []
    KD_list = []
    fam_list = []
    tp_list = []

    for d in range(len(timecourse)):
        tp = timecourse[d]
        freePan_no_na = freePan.sel(timepoint = tp).dropna("dim_0")
        cellnum = int(len(freePan_no_na)*samplefrac)
        if cellnum > 0:
            cell_id = np.random.choice(len(freePan_no_na), cellnum)
            cells = freePan_no_na[cell_id, :]
            kdl = list(cells.loc[dict(dim_1="affinity")].values)
            # transform norm E to KD
            kdll = np.exp(cf.y0+np.array(kdl)*cf.m)
            KD_list += list(kdll)
            SHM_list += list(list(cells.loc[dict(dim_1="mutations")].values))
            fam_list += list(cells.loc[dict(dim_1="family")].values)
            tp_list += [tp for k in range(cellnum)]

    # count into families and find clusters with more than xx members
    famcounter = Counter(fam_list)
    fams = list(famcounter.keys())
    clusters = []
    for fam in fams:
        if famcounter[fam] > 1:
            clusters.append(fam)

    # make separate lists for SHM, KD and TP (defining color) within clusters
    # and add information to list
    iSHMs = [[] for i in clusters]
    iKDs = [[] for i in clusters]
    iTPs = [[] for i in clusters]

    for ff in range(len(fam_list)):
        if fam_list[ff] in clusters:
            ii = clusters.index(fam_list[ff])
            iSHMs[ii].append(SHM_list[ff])
            iKDs[ii].append(KD_list[ff])
            # give different colors for different timepoints
            if tp_list[ff] == timecourse[0]:
                iTPs[ii].append('lightcoral')
            elif tp_list[ff] == timecourse[1]:
                iTPs[ii].append('indianred')
            else:
                iTPs[ii].append('firebrick')
    # pass information to plot function
    clonal_scatter_plot(iSHMs, iKDs, iTPs)

    # write information to file
    if d_export:
        datafile = open('processed_data/TUCHMI_sampling_data', 'w')

        datafile.write('1) SAMPLE STATISTICS \n \n')
        datafile.write('sampled fraction = {} \n \n'.format(samplefrac))
        datafile.write('timecourse (days) \n {} \n \n'.format(np.array(tList)/12.))
        datafile.write('SHMs of cells in sample, mean and std \n {} \n {} \n \n'.format(MSHM, SSHM))
        datafile.write('normalised Shannon entropy of cells in sample, mean and std \n {} \n {} \n \n'.format(MEntropies, SEntropies))
        datafile.write('fraction of non-unique cells in sample, mean and std \n {} \n {} \n \n'.format(Mclusterfracs, Sclusterfracs))

        datafile.close()

    return(simdata)

def selection_vs_mutation(store_export='dictionary', d_export=True):
    """ Performs a single simulation of a given size using a specified protocol
    of vaccination boosters. At specified timepoints, a specified number of
    memory cells is sampled and the affinities of their ancestors as well as
    their current affinities are written to a list. Also written to list
    are the binding energies of the naive cells. These three lists are then
    passed on to be plotted as distribution histograms.

    Plots produced include a collection of three histograms (unselected,
    selected germline energies, actual energies after mutations) for each
    queried timepoint and a more complex scatter plot with marginal histograms
    for each queried timepoint.

    For each timepoint, the fraction of cells with unaltered/improved/impaired
    affinity is printed to screen.

    If store_export is set 'datafile', the simulation data is stored in a
    hdf5 file for future purposes, for 'dictionary' the data is passed
    internally and lost after the run.

    If d_export is set True, textfiles containing the sampled data (used for
    plotting) are exported for each plot individually.
    """
    # parameters relevant to this analysis
    # evaluation timepoint in days
    analysis_times = [29]
    # prepare lists
    ancestor_dists = []
    final_dists = []
    # get runID from current system time
    runID = int(time.time())
    # run simulation and get filepath or dict
    simdata = main(runID, store_export=store_export, evalperday=1)
    # import required information for small scale plots
    l_times, l_fn, l_fm, l_GCs, LFcurve, Agcurve, evaltimes, freePan, GCPans, \
        ms_times, ms_vals, ms_fams, ms_muts, mut_list, E_list = \
        import_file(simdata)
    # extract the affinities and ancestor affinities at the analysis points
    tList = list(freePan["timepoint"].values)
    for i in range(len(analysis_times)):
        # limit cell number to be drawn in order not to clatter the plot
        tp = analysis_times[i]
        freePan_no_na = freePan.sel(timepoint = tList[tp]).dropna("dim_0")
        cellnum = min(2000, len(freePan_no_na))
        cellnum = len(freePan_no_na)
        cell_id = np.random.choice(len(freePan_no_na), cellnum)
        cells = freePan_no_na[cell_id, :]
        afflist = list(cells.loc[dict(dim_1="affinity")].values)
        final_dists.append(afflist)
        aff0list = list(cells.loc[dict(dim_1="affinity0")].values)
        ancestor_dists.append(aff0list)
    # send energy lists to histogram plot
    for i in range(len(analysis_times)):
        energy_distributions_plot(E_list, ancestor_dists[i], final_dists[i],
                                  analysis_times[i])
        energy_scatter_plot(ancestor_dists[i], final_dists[i],
                            analysis_times[i])

    if d_export:
        datafile = open('processed_data/energy_distribution_data', 'w')

        datafile.write('1) naive distribution \n \n')
        datafile.write('{} \n \n'.format(E_list))

        datafile.write('2) analysis days \n \n')
        datafile.write('{} \n \n'.format(analysis_times))

        datafile.write('3) ancestor distributions per time point \n \n')
        datafile.write('{} \n \n'.format(ancestor_dists))

        datafile.write('4) memory distributions per time point \n \n')
        datafile.write('{} \n \n'.format(final_dists))

        datafile.close()

    return(simdata)

def stacked_mutations(store_export='dictionary', d_export=True,
                      repeats=10):
    """ Performs a number of simulation runs, computes histograms for
    improved, impaired and unchanged binders at a single given timepoint
    and saves the individual as well as the summed values to file. For several
    repeats, data is accumulated in the histograms as well.

    If store_export is set 'datafile', the simulation data is stored in a
    hdf5 file for future purposes, for 'dictionary' the data is passed
    internally and lost after the run.

    If d_export is set True, textfiles containing the sampled data (used for
    plotting) are exported for each plot individually."""
    # parameters relevant to this analysis
    # evaluation timepoint in days
    analysis_time = 29
    bins = np.linspace(0.6, 1, 17)
    # collect results
    sum_zero = np.zeros(len(bins)-1)
    sum_plus = np.zeros(len(bins)-1)
    sum_minus = np.zeros(len(bins)-1)
    list_zero = []
    list_plus = []
    list_minus = []

    for i in range(repeats):
        # get runID from current system time
        runID = int(time.time())
        # run simulation and get filepath or dict
        simdata = main(runID, store_export=store_export, evalperday=1)
        # import required information for small scale plots
        l_times, l_fn, l_fm, l_GCs, LFcurve, Agcurve, evaltimes, freePan, \
            GCPans, ms_times, ms_vals, ms_fams, ms_muts, mut_list, E_list = \
            import_file(simdata)
        # extract the affinities and ancestor affinities at the analysis points
        tList = list(freePan["timepoint"].values)
        # limit cell number to be drawn in order not to clatter the plot
        # possibility of subsampling here
        tp = analysis_time

        freePan_no_na = freePan.sel(timepoint = tList[tp]).dropna("dim_0")
        cellnum = len(freePan_no_na)
        cell_id = np.random.choice(len(freePan_no_na), cellnum)
        cells = freePan_no_na[cell_id, :]

        final_dist = list(cells.loc[dict(dim_1="affinity")].dropna("dim_0").values)
        ancestor_dist = list(cells.loc[dict(dim_1="affinity0")].dropna("dim_0").values)

        # extract counts of unchanged, improved and impaired cells
        unchanged_list = np.array(final_dist)[np.where(np.array(
            ancestor_dist) == np.array(final_dist))[0]]
        improved_list = np.array(final_dist)[np.where(np.array(
            ancestor_dist) < np.array(final_dist))[0]]
        impaired_list = np.array(final_dist)[np.where(np.array(
            ancestor_dist) > np.array(final_dist))[0]]

        # make histograms, store information both in list and in sum.
        U_counts, _ = np.histogram(unchanged_list, bins=bins)
        plus_counts, _ = np.histogram(improved_list, bins=bins)
        minus_counts, _ = np.histogram(impaired_list, bins=bins)

        # collect results
        sum_zero += U_counts
        sum_plus += plus_counts
        sum_minus += minus_counts
        list_zero.append(U_counts)
        list_plus.append(plus_counts)
        list_minus.append(minus_counts)
    cellsum = np.sum(sum_zero)+np.sum(sum_plus)+np.sum(sum_minus)

    # plot
    stacked_energy_plot(bins, sum_plus, sum_minus, sum_zero, analysis_time)

    if d_export:
        datafile = open('processed_data/stacked_histogram_data', 'w')

        datafile.write('1) day \n \n')
        datafile.write('{} \n \n'.format(analysis_time))

        datafile.write('2) bins \n \n')
        datafile.write('{} \n \n'.format(bins))

        datafile.write('3) runs \n \n')
        datafile.write('{} \n \n'.format(repeats))

        datafile.write('4) sum of counts with unchanged energies\n \n')
        datafile.write('{} \n \n'.format(sum_zero))

        datafile.write('5) sum of counts with improved energies\n \n')
        datafile.write('{} \n \n'.format(sum_plus))

        datafile.write('6) sum of counts with impaired energies\n \n')
        datafile.write('{} \n \n'.format(sum_minus))

        datafile.write('7) list of counts with unchanged energies\n \n')
        datafile.write('{} \n \n'.format(list_zero))

        datafile.write('8) list of counts with improved energies\n \n')
        datafile.write('{} \n \n'.format(list_plus))

        datafile.write('9) list of counts with impaired energies\n \n')
        datafile.write('{} \n \n'.format(list_minus))

        datafile.write('10) percentage of umutated, improved, impaired \n \n')
        datafile.write('{}, {}, {}'.format(np.sum(sum_zero)/cellsum,
                       np.sum(sum_plus)/cellsum, np.sum(sum_minus)/cellsum))
        datafile.close()

def AM_effect_nkey(nkeys=[1, 5, 10, 15], repeats=100, d_export=True):
    """ Given a list of values for nkey and a number of individual GC reactions
    to be averaged over for each of them, computes and plots the improvement
    within single GCs for one infection. Thus, overwrites parameters giving
    the infection protocol and duration of the simulation as well as setting
    the nubmer of GCs to 1. Other parameters remain untouched. A textfile with
    the computed mean results is exported if d_export==True.
    """

    # set single infection and single GC for this analysis
    cf.endtime = 30*12
    cf.tinf = [0*12]
    cf.dose = [1]
    cf.nGCs = 1
    cf.naive_pool = 1000*1  # size of the naive precursor pool
    cf.memory_pool = 100*1  # size of the initial unspecific memory pool
    # function for calculating mean E_norm from GC panel

    def GC_affinity(GCPan):
        """ Given a GC panel, gets the mean E_norm for each timepoint."""
        energies = []
        tList = list(GCPan["timepoint"].values)

        for tp in range(len(tList)):
            energy = GCPan.sel(timepoint = tList[tp]).loc[dict(dim_1="affinity")].dropna("dim_0").values.mean()
            energies.append(energy)

        return tList, energies

    topElist = []
    for hs in nkeys:
        # set binding model parameters accordingly
        cf.nkey = hs
        cf.lAg = hs
        cf.lAb = 220 - hs
        eL = []  # list for collecting energies timecurses of all runs
        for r in range(repeats):
            simdata = main(store_export='dictionary', evalperday=12)
            l_times, l_fn, l_fm, l_GCs, LFcurve, Agcurve, evaltimes, freePan, \
                GCPans, ms_times, ms_vals, ms_fams, ms_muts, mut_list, E_list\
                = import_file(simdata)
            tList, energies = GC_affinity(GCPans[0])
            eL.append(energies)

        # calculate mean and std of all runs and plot
        eM = np.nanmean(np.array(eL), axis=0)
        eStd = np.nanstd(np.array(eL), axis=0)
        topElist.append((hs, tList, eM, eStd))

    # write information to file
    if d_export:
        datafile = open('processed_data/AM_effect_data', 'w')
        datafile.write('number of simulation runs per n_key = {} \n'.format(repeats))
        datafile.write('n_key, time (days), mean(normalised energies), std(normalised energies) \n')
        for i in range(len(nkeys)):
            datafile.write('{0}, {1}, {2}, {3}\n \n'.format(topElist[i][0], np.array(topElist[i][1])/12., topElist[i][2], topElist[i][3]))
        datafile.close()
    # plot
    AM_effect_plot(topElist)


"""
By Dr Jie Zheng -Q, NAOC
v1 2019-04-27
"""

import numpy as np
from..util import *

def date_conv():
    pass

  #function date_conv,date,type, BAD_DATE = bad_date
  #;+
  #; NAME:
  #;     DATE_CONV
  #; PURPOSE:
  #;     Procedure to perform conversion of dates to one of three possible formats.
  #;
  #; EXPLANATION:
  #;     The following date formats are allowed
  #;
  #;       format 1: real*8 scalar encoded as:
  #;               year*1000 + day + hour/24. + min/24./60 + sec/24./60/60
  #;               where day is the day of year (1 to 366)
  #;       format 2: Vector encoded as:
  #;               date[0] = year (eg. 2005)
  #;               date[1] = day of year (1 to 366)
  #;               date[2] = hour
  #;               date[3] = minute
  #;               date[4] = second
  #;               To indicate a date only, set a negative hour.
  #;       format 3: string (ascii text) encoded as
  #;               DD-MON-YEAR HH:MM:SS.SS
  #;               (eg.  14-JUL-2005 15:25:44.23)
  #;            OR
  #;               YYYY-MM-DD HH:MM:SS.SS  (ISO standard)
  #;               (eg.  1987-07-14 15:25:44.23 or 1987-07-14T15:25:44.23)
  #;
  #;            OR
  #;               DD/MM/YY (pre-2000 option for FITS DATE keywords)
  #;            Time of day segment is optional in all of these.
  #;
  #;       format 4: three element vector giving spacecraft time words
  #;       from a Hubble Space Telescope (HST) telemetry packet.   Based on
  #;       total number of secs since midnight, JAN. 1, 1979
  #;
  #;       format 5: Julian day. As this is also a scalar, like format 1,
  #;       	the distinction between the two on input is made based on their
  #;       	value. Numbers > 2300000 are interpreted as Julian days.
  #;
  #; CALLING SEQUENCE
  #;       results = DATE_CONV( DATE, TYPE )
  #;
  #; INPUTS:
  #;       DATE - input date in one of the possible formats. Must be scalar.
  #;       TYPE - type of output format desired.  If not supplied then
  #;               format 3 (real*8 scalar) is used.
  #;                       valid values:
  #;                       'REAL'  - format 1
  #;                       'VECTOR' - format 2
  #;                       'STRING' - format 3
  #;                       'FITS' - YYYY-MM-DDTHH:MM:SS.SS'
  #;                       'JULIAN' - Julian date
  #;                       'MODIFIED' - Modified Julian date (JD-2400000.5)
  #;               TYPE can be abbreviated to the single character strings 'R',
  #;               'V', 'S', 'F', 'J', and 'M'.
  #;               Nobody wants to convert TO spacecraft time (I hope!)
  #; OUTPUTS:
  #;       The converted date is returned as the function value.
  #;       Output is -1 if date is unrecognisable.
  #;
  #;       If the time of day is omitted from the input, it will also
  #;       be omitted from any output string (format STRING or FITS).
  #;       Note that date-only strings are allowed by the FITS standard.
  #;       For other output formats any missing time of day is set to
  #;       00:00:00.0
  #;
  #; KEYWORD OUTPUTS
  #;
  #;        BAD_DATE set to 1B if date is unrecognisable
  #;
  #; EXAMPLES:
  #;       IDL> print,date_conv('2006-03-13 19:58:00.00'),f='(f15.5)'
  #;             2006072.83194
  #;       IDL> print,date_conv( 2006072.8319444d,'F')
  #;             2006-03-13T19:58:00.00
  #;       IDL> print,date_conv( 2006072.8319444d,'V')
  #;             2006.00      72.0000      19.0000      57.0000      59.9962
  #;       IDL> print,date_conv( 2006072.8319444d,'J'), f='(f15.5)'
  #;             2453808.33194
  #;
  #;
  #; HISTORY:
  #;      version 1  D. Lindler  July, 1987
  #;      adapted for IDL version 2  J. Isensee  May, 1990
  #;      Made year 2000 compliant; allow ISO format input  jls/acc Oct 1998
  #;      DJL/ACC Jan 1998, Modified to work with dates such as 6-JAN-1996 where
  #;               day of month has only one digit.
  #;      DJL, Nov. 2000, Added input/output format YYYY-MM-DDTHH:MM:SS.SS
  #;      Replace spaces with '0' in output FITS format  W.Landsman April 2006
  #;      Added Julian date capabilities on input and output.  M.Perrin, July 2007
  #;      Removed spurious /WARN keyword to MESSAGE W.L. Feb 2012
  #;      ...and another /WARN; added BAD_DATE, drop spurious time-of-day
  #;      output from strings. J. P. Leahy July 2013
  #;      changed all /CONTINUE warning messages to /INFO: can be suppressed
  #;      by setting !QUIET = 1.  J. P. Leahy July 2013
  #;-
  #;-------------------------------------------------------------
  #;
  #compile_opt idl2
  #; data declaration
  #;
  #days = [0,31,28,31,30,31,30,31,31,30,31,30,31]
  #months = ['   ','JAN','FEB','MAR','APR','MAY','JUN','JUL','AUG','SEP','OCT',$
  #        'NOV','DEC']
  #;
  #; set default type if not supplied
  #;
  #if N_params() lt 2 then type = 'REAL'
  #;
  #; Determine type of input supplied
  #;
  #s = size(date) & ndim = s[0] & datatype = s[ndim+1]
  #if ndim gt 0 then begin                 ;vector?
  #        if ndim gt 1 then goto,notvalid
  #        if (s[1] ne 5) && (s[1] ne 3) then goto,notvalid
  #        if (s[1] eq 5) then form = 2 else form = 4
  #   end else begin                       ;scalar input
  #        if datatype eq 0 then goto,notvalid
  #        if datatype eq 7 then form = 3 $        ;string
  #                         else form = 1  ;numeric scalar
  #end
  #;
  #;      -----------------------------------
  #;
  #;*** convert input to year,day,hour,minute,second
  #;
  #;      -----------------------------------
  #case form of
  #
  #        1: begin                                        ;real scalar
  #			; The 'real' input format may be interpreted EITHER
  #			; a) if < 2300000
  #			;    as the traditional 'real*8 encoded' format used by date_conv
  #			; b) if > 2300000
  #			;    as a Julian Day Number
  #                idate = long(date)
  #                year = long(idate/1000)
  #
  #				if year lt 2300 then begin
  #
  #					; if year is only 2 digits, assume 1900
  #	                if year lt 100 then begin
  #	                   message,/INF, $
  #	                     'Warning: Year specified is only 2 digits, assuming 19xx'
  #	                   year=1900+year
  #	                   idate=1900000+idate
  #	                   date=1900000.+date
  #	                end
  #	                day = idate - year*1000
  #	                fdate = date-idate
  #	                fdate = fdate*24.
  #	                hour = fix(fdate)
  #	                fdate = (fdate-hour)*60.0
  #	                minute = fix(fdate)
  #	                sec = float((fdate-minute)*60.0)
  #
  #				endif else begin
  #					daycnv, date, year, mn, mndy, hr
  #					; convert from month/day to day of year
  #					; how many days PRECEED the start of each month?
  #					YDAYS = [0,31,59,90,120,151,181,212,243,273,304,334,366]
  #					LEAP =  (((YeaR MOD 4) EQ 0) AND ((YeaR MOD 100) NE 0)) OR $
  #					                 ((YeaR MOD 400) EQ 0)
  #			        IF LEAP THEN YDAYS[2:*] = YDAYS[2:*] + 1
  #					day = ydays[mn-1]+mndy
  #
  #					hour = fix(hr)
  #					fmin = (hr-hour)*60
  #					minute = fix(fmin)
  #					sec = float((fmin-minute)*60)
  #				endelse
  #           end
  #
  #        2: begin                                        ;vector
  #                year = fix(date[0])
  #;
  #; if year is only 2 digits, assume 1900
  #;
  #                if year lt 100 then begin
  #                   message,/INF, $
  #                    'Warning: Year specified is only 2 digits, assuming 19xx'
  #                   year=1900+year
  #                end
  #;
  #                day = fix(date[1])
  #                hour = fix(date[2])
  #                minute = fix(date[3])
  #                sec = float(date[4])
  #           end
  #
  #        3: begin                                        ;string
  #                temp = date
  #;
  #; check for old type of date, DD-MMM-YYYY
  #;
  #                test = STRPOS(temp,'-')
  #                if test ge 0 && test le 2 then begin
  #                  day_of_month = fix(gettok(temp,'-'))
  #                  month_name = gettok(temp,'-')
  #                  year = fix(gettok(temp,' '))
  #;
  #; determine month number from month name
  #;
  #                  month_name = strupcase(month_name)
  #                  for mon = 1,12 do begin
  #                        if month_name eq months[mon] then goto,found
  #                  end
  #                  message,/INFORMATIONAL, 'Invalid month name specified'
  #                  goto, notvalid
  #;
  #; check for new type of date, ISO: YYYY-MM-DD
  #;
  #                end else if strpos(temp,'-') eq 4 then begin
  #                  year = fix(gettok(temp,'-'))
  #                  month_name = gettok(temp,'-')
  #                  mon= FIX(month_name)
  #                  day_of_month=gettok(temp,' ')
  #                  if strlen(temp) eq 0 then begin
  #                        dtmp=gettok(day_of_month,'T')
  #                        temp=day_of_month
  #                        day_of_month=dtmp
  #                  end
  #                  day_of_month=fix(day_of_month)
  #;
  #; check for DD/MM/YY
  #;
  #                end else if STRPOS(temp,'/') eq 2 then begin
  #                  day_of_month = FIX(gettok(temp,'/'))
  #                  mon = FIX(gettok(temp,'/'))
  #                  year = 1900 + FIX(STRMID(temp,0,2))
  #                end else goto, notvalid
  #
  #    found:
  #                hour = gettok(temp,':')
  #                hour =  hour NE '' ? FIX(hour) : -1
  #                minute = fix(gettok(temp,':'))
  #                sec = float(strtrim(strmid(temp,0,5)))
  #
  #                IF (mon LT 1 || mon GT 12) THEN BEGIN
  #                    MESSAGE, /INFORMATIONAL, 'Invalid month specified'
  #                    goto, notvalid
  #                ENDIF
  #;
  #; if year is only 2 digits, assume 1900
  #;
  #                if year lt 100 then begin
  #                   message,/INFORMATIONAL, $
  #                     'Warning: Year specified is only 2 digits, assuming 19xx'
  #                   year=1900+year
  #                end
  #;
  #;
  #;            convert to day of year from month/day_of_month
  #;
  #;            correction for leap years
  #;
  #;               if (fix(year) mod 4) eq 0 then days(2) = 29     ;add one to february
  #                lpyr = ((year mod 4) eq 0) and ((year mod 100) ne 0) $
  #                        or ((year mod 400) eq 0)
  #                if lpyr eq 1 then days[2] = 29 ; if leap year, add day to Feb.
  #;
  #;
  #;            compute day of year
  #;
  #                  day = fix(total(days[0:mon-1])+day_of_month)
  #           end
  #
  #        4 : begin                       ;spacecraft time
  #                SC = DOUBLE(date)
  #                SC = SC + (SC LT 0.0)*65536.    ;Get rid of neg. numbers
  #;
  #;            Determine total number of secs since midnight, JAN. 1, 1979
  #;
  #                SECS = SC[2]/64 + SC[1]*1024 + SC[0]*1024*65536.
  #                SECS = SECS/8192.0D0            ;Convert from spacecraft units
  #;
  #;            Determine number of years
  #;
  #                MINS = SECS/60.
  #                HOURS = MINS/60.
  #                TOTDAYS = HOURS/24.
  #                YEARS = TOTDAYS/365.
  #                YEARS = FIX(YEARS)
  #;
  #;            Compute number of leap years past
  #;
  #                LEAPYEARS = (YEARS+2)/4
  #;
  #;           Compute day of year
  #;
  #                DAY = FIX(TOTDAYS-YEARS*365.-LEAPYEARS)
  #;
  #;           Correct for case of being right at end of leapyear
  #;
  #                IF DAY LT 0 THEN BEGIN
  #                  DAY = DAY+366
  #                  LEAPYEARS = LEAPYEARS-1
  #                  YEARS = YEARS-1
  #                END
  #;
  #;            COMPUTE HOUR OF DAY
  #;
  #                TOTDAYS = YEARS*365.+DAY+LEAPYEARS
  #                HOUR = FIX(HOURS - 24*TOTDAYS)
  #                TOTHOURS = TOTDAYS*24+HOUR
  #;
  #;            COMPUTE MINUTE
  #;
  #                MINUTE = FIX(MINS-TOTHOURS*60)
  #                TOTMIN = TOTHOURS*60+MINUTE
  #;
  #;            COMPUTE SEC
  #;
  #                SEC = SECS-TOTMIN*60
  #;
  #;            COMPUTE ACTUAL YEAR
  #;
  #                YEAR = YEARS+79
  #;
  #; if year is only 2 digits, assume 1900
  #;
  #                if year lt 100 then begin
  #                   message, /INF, $
  #                     'Warning: Year specified is only 2 digits, assuming 19xx'
  #                   year=1900+year
  #                end
  #;
  #;
  #;            START DAY AT ONE AND NOT ZERO
  #;
  #                DAY++
  #           END
  #ENDCASE
  #;
  #;            correction for leap years
  #;
  #        if form ne 3 then begin         ;Was it already done?
  #           lpyr = ((year mod 4) eq 0) && ((year mod 100) ne 0) $
  #                || ((year mod 400) eq 0)
  #           if lpyr eq 1 then days[2] = 29 ; if leap year, add day to Feb.
  #        end
  #;
  #;            check for valid day
  #;
  #        if (day lt 1) || (day gt total(days)) then begin
  #            message, /INFORMATIONAL, $
  #                'ERROR -- There are only ' + strtrim(fix(total(days)),2) + $
  #	              ' days  in year '+strtrim(year,2)
  #            goto, notvalid
  #        endif
  #;
  #;            find month which day occurs
  #;
  #        day_of_month = day
  #        month_num = 1
  #        while day_of_month gt days[month_num] do begin
  #               day_of_month = day_of_month - days[month_num]
  #               month_num = month_num+1
  #        end
  #;           ---------------------------------------
  #;
  #;   *****       Now convert to output format
  #;
  #;           ---------------------------------------
  #;
  #; is type a string
  #;
  #s = size(type)
  #if (s[0] ne 0) or (s[1] ne 7) then $
  #        message,'ERROR - Output type specification must be a string'
  #;
  #outcode = STRMID(STRUPCASE(type),0,1)
  #IF (outcode EQ 'S' || outcode EQ 'F') && hour GE 0 THEN BEGIN
  #    xsec = strmid(string(sec+100,'(f6.2)'),1,5)
  #    if xsec EQ '60.00' then begin
  #        minute = minute+1
  #        xsec = '00.00'
  #    endif
  #    xminute =   string(minute,'(i2.2)')
  #    if xminute EQ '60' then begin
  #        hour = hour+1
  #        xminute = '00'
  #    endif
  #    tod = string(hour,'(i2.2)') +  ':' +xminute + ':'+ xsec
  #ENDIF
  #
  #case outcode of
  #
  #        'V' : begin                             ;vector output
  #                out = fltarr(5)
  #                out[0] = year
  #                out[1] = day
  #                out[2] = hour > 0
  #                out[3] = minute
  #                out[4] = sec
  #             end
  #
  #        'R' : begin                             ;floating point scalar
  #;               if year gt 1900 then year = year-1900
  #                out = sec/24.0d0/60./60. + minute/24.0d0/60. $
  #                + (hour > 0)/24.0d0  +  day + year*1000d0
  #              end
  #
  #        'S' : begin                             ;string output
  #
  #                month_name = months[month_num]
  #;
  #;            encode into ascii_date
  #;
  #                out = string(day_of_month,'(i2)') +'-'+ month_name +'-' + $
  #                        string(year,'(i4)')
  #
  #  ; Omit time of day from output string if not specified on input
  #                IF hour GE 0 THEN out += ' '+tod
  #           end
  #        'F' : begin
  #                out = string(year,'(i4)')+'-'+string(month_num,'(I2.2)') $
  #                      + '-' +  string(day_of_month,'(i2.2)')
  #                IF hour GE 0 THEN out += 'T' + tod
  #           end
  #
  #		'J' : begin	; Julian Date
  #				ydn2md, year, day, mn, dy
  #				juldate, [year, mn, dy, hour, minute, sec], rjd
  #				out = rjd+2400000   ; convert from reduced to regular JD
  #			  end
  #		'M' : begin ; Modified Julian Date = JD - 2400000.5
  #				ydn2md, year, day, mn, dy
  #				juldate, [year, mn, dy, hour, minute, sec], rjd
  #				out = rjd-0.5   ; convert from reduced to modified JD
  #			  end
  #
  #        else: begin                     ;invalid type specified
  #                print,'DATE_CONV-- Invalid output type specified'
  #                print,' It must be ''REAL'', ''STRING'', ''VECTOR'', ''JULIAN'', ''MODIFIED'', or ''FITS''.'
  #                return,-1
  #              end
  #endcase
  #
  #bad_date = 0B
  #return,out
  #;
  #; invalid input date error section
  #;
  #NOTVALID:
  #bad_date = 1B
  #message, 'Invalid input date specified', /INFORMATIONAL
  #return, -1
  #end


# Copyright 2022 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""target assigner"""
import logging

import numpy as np
import numpy.random as npr

logger = logging.getLogger(__name__)

def unmap(data, count, inds, fill=0):
    """
    Unmap a subset of item (data) back to the original set of items
    (of size count)
    """
    if count == len(inds):
        return data

    if len(data.shape) == 1:
        ret = np.empty((count,), dtype=data.dtype)
        ret.fill(fill)
        ret[inds] = data
    else:
        ret = np.empty((count,) + data.shape[1:], dtype=data.dtype)
        ret.fill(fill)
        ret[inds, :] = data
    return ret

def create_target_np(all_anchors,
                     gt_boxes,
                     similarity_fn,
                     box_encoding_fn,
                     prune_anchor_fn=None,
                     gt_classes=None,
                     matched_threshold=0.6,
                     unmatched_threshold=0.45,
                     positive_fraction=None,
                     rpn_batch_size=300,
                     norm_by_num_examples=False,
                     box_code_size=7):
    """Modified from FAIR detectron.
    Args:
        all_anchors: [num_of_anchors, box_ndim] float tensor.
        gt_boxes: [num_gt_boxes, box_ndim] float tensor.
        similarity_fn: a function, accept anchors and gt_boxes, return
            similarity matrix(such as IoU).
        box_encoding_fn: a function, accept gt_boxes and anchors, return
            box encodings(offsets).
        prune_anchor_fn: a function, accept anchors, return indices that
            indicate valid anchors.
        gt_classes: [num_gt_boxes] int tensor. indicate gt classes, must
            start with 1.
        matched_threshold: float, iou greater than matched_threshold will
            be treated as positives.
        unmatched_threshold: float, iou smaller than unmatched_threshold will
            be treated as negatives.
        positive_fraction: [0-1] float or None. if not None, we will try to
            keep ratio of pos/neg equal to positive_fraction when sample.
            if there is not enough positives, it fills the rest with negatives
        rpn_batch_size: int. sample size
        norm_by_num_examples: bool. norm box_weight by number of examples, but
            I recommend to do this outside.
        box_code_size: int. box coder size
    Returns:
        labels, bbox_targets, bbox_outside_weights
    """
    total_anchors = all_anchors.shape[0]
    if prune_anchor_fn is not None:
        inds_inside = prune_anchor_fn(all_anchors)
        anchors = all_anchors[inds_inside, :]
        if not isinstance(matched_threshold, float):
            matched_threshold = matched_threshold[inds_inside]
        if not isinstance(unmatched_threshold, float):
            unmatched_threshold = unmatched_threshold[inds_inside]
    else:
        anchors = all_anchors
        inds_inside = None
    num_inside = len(inds_inside) if inds_inside is not None else total_anchors

    if gt_classes is None:
        gt_classes = np.ones([gt_boxes.shape[0]], dtype=np.int32)
    labels = np.empty((num_inside,), dtype=np.int32)
    gt_ids = np.empty((num_inside,), dtype=np.int32)
    labels.fill(-1)
    gt_ids.fill(-1)
    if np.array(gt_boxes).shape[0] > 0 and anchors.shape[0] > 0:
        anchor_by_gt_overlap = similarity_fn(anchors, gt_boxes)
        anchor_to_gt_argmax = anchor_by_gt_overlap.argmax(axis=1)
        anchor_to_gt_max = anchor_by_gt_overlap[np.arange(num_inside), anchor_to_gt_argmax]
        gt_to_anchor_argmax = anchor_by_gt_overlap.argmax(axis=0)
        gt_to_anchor_max = anchor_by_gt_overlap[gt_to_anchor_argmax,
                                                np.arange(anchor_by_gt_overlap.shape[1])]
        empty_gt_mask = gt_to_anchor_max == 0
        gt_to_anchor_max[empty_gt_mask] = -1
        anchors_with_max_overlap = np.where(anchor_by_gt_overlap == gt_to_anchor_max)[0]
        gt_inds_force = anchor_to_gt_argmax[anchors_with_max_overlap]
        labels[anchors_with_max_overlap] = gt_classes[gt_inds_force]
        gt_ids[anchors_with_max_overlap] = gt_inds_force
        pos_inds = anchor_to_gt_max >= matched_threshold
        gt_inds = anchor_to_gt_argmax[pos_inds]
        labels[pos_inds] = gt_classes[gt_inds]
        gt_ids[pos_inds] = gt_inds
        bg_inds = np.where(anchor_to_gt_max < unmatched_threshold)[0]
    else:
        bg_inds = np.arange(num_inside)
    fg_inds = np.where(labels > 0)[0]
    fg_max_overlap = None
    if np.array(gt_boxes).shape[0] > 0 and anchors.shape[0] > 0:
        fg_max_overlap = anchor_to_gt_max[fg_inds]
    gt_pos_ids = gt_ids[fg_inds]
    if positive_fraction is not None:
        num_fg = int(positive_fraction * rpn_batch_size)
        if len(fg_inds) > num_fg:
            disable_inds = npr.choice(fg_inds, size=(len(fg_inds) - num_fg), replace=False)
            labels[disable_inds] = -1
            fg_inds = np.where(labels > 0)[0]

        num_bg = rpn_batch_size - np.sum(labels > 0)
        if len(bg_inds) > num_bg:
            enable_inds = bg_inds[npr.randint(len(bg_inds), size=num_bg)]
            labels[enable_inds] = 0
    else:
        if np.array(gt_boxes).shape[0] == 0 or anchors.shape[0] == 0:
            labels[:] = 0
        else:
            labels[bg_inds] = 0
            labels[anchors_with_max_overlap] = gt_classes[gt_inds_force]
    bbox_targets = np.zeros((num_inside, box_code_size), dtype=all_anchors.dtype)
    if np.array(gt_boxes).shape[0] > 0 and anchors.shape[0] > 0:
        bbox_targets[fg_inds, :] = box_encoding_fn(gt_boxes[anchor_to_gt_argmax[fg_inds], :],
                                                   anchors[fg_inds, :])

    bbox_outside_weights = np.zeros((num_inside,), dtype=all_anchors.dtype)
    if norm_by_num_examples:
        num_examples = np.sum(labels >= 0)  # neg + pos
        num_examples = np.maximum(1.0, num_examples)
        bbox_outside_weights[labels > 0] = 1.0 / num_examples
    else:
        bbox_outside_weights[labels > 0] = 1.0

    if inds_inside is not None:
        labels = unmap(labels, total_anchors, inds_inside, fill=-1)
        bbox_targets = unmap(bbox_targets, total_anchors, inds_inside, fill=0)
        bbox_outside_weights = unmap(bbox_outside_weights, total_anchors, inds_inside, fill=0)
    ret = {
        "labels": labels,
        "bbox_targets": bbox_targets,
        "bbox_outside_weights": bbox_outside_weights,
        "assigned_anchors_overlap": fg_max_overlap,
        "positive_gt_id": gt_pos_ids,
    }
    if inds_inside is not None:
        ret["assigned_anchors_inds"] = inds_inside[fg_inds]
    else:
        ret["assigned_anchors_inds"] = fg_inds
    return ret

class TargetAssigner:
    """target assigner"""

    def __init__(self,
                 box_coder,
                 anchor_generators,
                 region_similarity_calculator=None,
                 positive_fraction=None,
                 sample_size=512):
        self._region_similarity_calculator = region_similarity_calculator
        self._box_coder = box_coder
        self._anchor_generators = anchor_generators
        self._positive_fraction = positive_fraction
        self._sample_size = sample_size

    @property
    def box_coder(self):
        """box coder"""
        return self._box_coder

    def assign(self,
               anchors,
               gt_boxes,
               anchors_mask=None,
               gt_classes=None,
               matched_thresholds=None,
               unmatched_thresholds=None):
        """assign"""

        def similarity_fn(anchors, gt_boxes):
            """similarity fn"""
            anchors_rbv = anchors[:, [0, 1, 3, 4, 6]]
            gt_boxes_rbv = gt_boxes[:, [0, 1, 3, 4, 6]]
            return self._region_similarity_calculator.compare(anchors_rbv, gt_boxes_rbv)

        def box_encoding_fn(boxes, anchors):
            """box encoding fn"""
            return self._box_coder.encode(boxes, anchors)

        if anchors_mask is not None:
            prune_anchor_fn = lambda _: np.where(anchors_mask)[0]
        else:
            prune_anchor_fn = None

        return create_target_np(anchors,
                                gt_boxes,
                                similarity_fn,
                                box_encoding_fn,
                                prune_anchor_fn=prune_anchor_fn,
                                gt_classes=gt_classes,
                                matched_threshold=matched_thresholds,
                                unmatched_threshold=unmatched_thresholds,
                                positive_fraction=self._positive_fraction,
                                rpn_batch_size=self._sample_size,
                                norm_by_num_examples=False,
                                box_code_size=self.box_coder.code_size)

    def generate_anchors(self, feature_map_size):
        """generate anchors"""
        anchors_list = []
        matched_thresholds = [
            a.match_threshold for a in self._anchor_generators
        ]
        unmatched_thresholds = [
            a.unmatch_threshold for a in self._anchor_generators
        ]
        match_list, unmatch_list = [], []
        for anchor_generator, match_thresh, unmatch_thresh in zip(self._anchor_generators,
                                                                  matched_thresholds,
                                                                  unmatched_thresholds):
            anchors = anchor_generator.generate(feature_map_size)
            anchors = anchors.reshape([*anchors.shape[:3], -1, 7])
            anchors_list.append(anchors)
            num_anchors = np.prod(anchors.shape[:-1])
            match_list.append(np.full([num_anchors], match_thresh, anchors.dtype))
            unmatch_list.append(np.full([num_anchors], unmatch_thresh, anchors.dtype))
        anchors = np.concatenate(anchors_list, axis=-2)
        matched_thresholds = np.concatenate(match_list, axis=0)
        unmatched_thresholds = np.concatenate(unmatch_list, axis=0)
        return {
            "anchors": anchors,
            "matched_thresholds": matched_thresholds,
            "unmatched_thresholds": unmatched_thresholds
        }

    @property
    def num_anchors_per_location(self):
        """num anchors per location"""
        num = 0
        for a_generator in self._anchor_generators:
            num += a_generator.num_anchors_per_localization
        return num


import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

import tensorflow as tf
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.compat.v1.Session(config=config)
import matplotlib.pyplot as plt
import tensorflow_probability as tfp
tfd = tfp.distributions

from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

import argparse

parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('--jobid', type=int, default=0, help='an integer for the accumulator')
parser.add_argument('--eps', type=float, default=0.001, help='step size')
parser.add_argument('--seed', type=int, default=100, help='random seed')
parser.add_argument('--nc', type=int, default=32, help='mesh size')
parser.add_argument('--bs', type=float, default=200, help='box size')
parser.add_argument('--suffix', type=str, default="", help='suffix to fpath')
parser.add_argument('--nR', type=int, default=0, help="number of smoothings")
parser.add_argument('--scaleprior', type=int, default=1, help="add power to scale to prior")
parser.add_argument('--nchains', type=int, default=1, help="number of chains")
parser.add_argument('--debug', type=int, default=0, help="debug run")
parser.add_argument('--reconiter', type=int, default=100, help="number of iterations for reconstruction")
parser.add_argument('--burnin', type=int, default=200, help="number of iterations for burnin")
parser.add_argument('--tadapt', type=int, default=50, help="number of iterations for eps adaptation")
parser.add_argument('--ntrain', type=int, default=10, help="number of training iterations")
parser.add_argument('--prentrain', type=int, default=200, help="number of training iterations")
parser.add_argument('--thinning', type=int, default=10, help="thinning")
parser.add_argument('--kwts', type=int, default=4, help='number of trenf layers')
parser.add_argument('--lpsteps1', type=int, default=25, help="min leapfrog steps")
parser.add_argument('--lpsteps2', type=int, default=50, help="max leapfrog steps")
parser.add_argument('--mcmciter', type=int, default=5, help="number of only mcmc iterations")
parser.add_argument('--order', type=int, default=1, help="ZA or LPT")
#
parser.add_argument('--mode', type=str, default="classic", help='mode for trenf')
parser.add_argument('--nlayers', type=int, default=3, help='number of trenf layers')
parser.add_argument('--nbins', type=int, default=32, help="number of bins in trenf spline")
parser.add_argument('--fitnoise', type=int, default=1, help='fitnoise')
parser.add_argument('--fitscale', type=int, default=1, help='fitscale')
parser.add_argument('--fitmean', type=int, default=1, help='fitmean')
parser.add_argument('--meanfield', type=int, default=1, help='meanfield for affine')
parser.add_argument('--nknots', type=int, default=100, help='number of trenf layers')
parser.add_argument('--linknots', type=int, default=0, help='linear spacing for knots')
#
parser.add_argument('--probjump', type=float, default=0.3, help='probability of jump')
parser.add_argument('--regwt0', type=float, default=0., help='regularization weight')

args = parser.parse_args()
device = args.jobid
suffix = args.suffix

from tensorflow.python.client import device_lib

def get_available_gpus():
    local_device_protos = device_lib.list_local_devices()
    return [x.name for x in local_device_protos if x.device_type == 'GPU']
#print("\nDevice name\n", tf.test.gpu_device_name(), "\n")
print("\nDevices\n", get_available_gpus())

import sys, os, time
import flowpm
sys.path.append('../../galference/utils/')
import tools
import diagnostics as dg

sys.path.append('../src/')
from pmfuncs import Evolve
from pyhmc import PyHMC, PyHMC_batch, DualAveragingStepSize
from trenfmodel import VItrenf
from hmcfuncs import DM_config, DM_fourier, Kwts
import recon, trenfmodel
from callback import callback, datafig, callback_fvi, callback_sampling

##########
nchains = args.nchains
reconiter = args.reconiter
burnin = args.burnin
mcmciter = args.mcmciter
ntrain = args.ntrain
tadapt = args.tadapt
thinning = args.thinning
lpsteps1, lpsteps2 = args.lpsteps1, args.lpsteps2
nsamples = 1
prentrain = args.prentrain
nR = args.nR
regwt0 = args.regwt0
recalib = 1
#
if args.debug:
    reconiter = 5
    burnin = 5
    mcmciter = 5
    tadapt = 5
    thinning = 2
    lpsteps1, lpsteps2 = 3, 5
    nsamples = 1
    prentrain = 10
    nR = 0
probjump = args.probjump
allRs = [0., 1., 2., 4.]
allR = allRs[:nR + 1][::-1]

#
bs, nc = args.bs, args.nc
nsteps = 3
a0, af, nsteps = 0.1, 1.0,  nsteps
stages = np.linspace(a0, af, nsteps, endpoint=True)
nsims = 200
donbody = False
order = args.order
shotnoise = bs**3/nc**3
dnoise = 1. #shotnoise/nc**1.5
if donbody: fpath = '/mnt/ceph/users/cmodi/galference/dm_fvi/L%04d_N%04d_T%02d'%(bs, nc, nsteps)
elif order == 2: fpath = '/mnt/ceph/users/cmodi/galference/dm_fvi/L%04d_N%04d_LPT'%(bs, nc)
elif order == 1: fpath = '/mnt/ceph/users/cmodi/galference/dm_fvi/L%04d_N%04d_ZA'%(bs, nc)
if suffix == "": fpath = fpath + '/'
else: fpath = fpath + "-" + suffix + '/'
os.makedirs('%s'%fpath, exist_ok=True)
os.makedirs('%s'%fpath + '/figs/', exist_ok=True)
os.makedirs('%s'%fpath + '/opt/', exist_ok=True)
os.makedirs('%s'%fpath + '/weights/', exist_ok=True)
os.makedirs('%s'%fpath + '/burnin/', exist_ok=True)

evolve = Evolve(nc, bs, a0=a0, af=af, nsteps = nsteps, donbody=donbody, order=order)

##############################################
##Generate DATA
print("\nFor seed : ", args.seed)
np.random.seed(args.seed)
zic = np.random.normal(0, 1, nc**3).reshape(1, nc, nc, nc).astype(np.float32)
noise = np.random.normal(0, dnoise, nc**3).reshape(1, nc, nc, nc).astype(np.float32)
ic = evolve.z_to_lin(zic).numpy()
fin = evolve.pm(tf.constant(ic)).numpy()
data = fin + noise
data = data.astype(np.float32)
tfdata = tf.constant(data)
tfnoise = tf.constant(dnoise)
np.save(fpath + 'ic', ic)
np.save(fpath + 'fin', fin)
np.save(fpath + 'data', data)
fig = datafig(ic, fin, data, bs, dnoise)
plt.savefig(fpath + 'data')
plt.close()
k, pic = tools.power(ic[0], boxsize=bs)
k, pf = tools.power(fin[0], boxsize=bs)
k, pd = tools.power(data[0], boxsize=bs)
k, pn = tools.power(1+noise[0], boxsize=bs)
knoise = evolve.kmesh[evolve.kmesh > k[(pn > pf)][0]].min()
print("Noise dominated after : ", knoise, (evolve.kmesh > knoise).sum()/nc**3)

dmfuncs = DM_fourier(evolve, tfdata, dnoise=dnoise)
kwts = Kwts(evolve, mode=args.kwts, knoise=knoise)
py_log_prob = lambda x: dmfuncs.unnormalized_log_prob(tf.constant(x, dtype=tf.float32), tf.constant(1.)).numpy().astype(np.float32)
py_grad_log_prob = lambda x: dmfuncs.grad_log_prob(tf.constant(x, dtype=tf.float32), tf.constant(1.)).numpy().astype(np.float32)
hmckernel = PyHMC_batch(py_log_prob, py_grad_log_prob, invmetric_diag=kwts.kwts, returnV=True)
epsadapt = DualAveragingStepSize(args.eps)
stepsize = args.eps
samples, pyacc = [], []

print("HMC kernels setup")
##############################################

def flow_step(model, x, lpx):
    nsamples = x.shape[0]
    xk = tf.identity(x)
    x = evolve.zk_to_z(x)
    lqx = model.q.log_prob(x)

    y = model.q.sample(nsamples)*1.
    yk = evolve.z_to_zk(y)
    lpy = dmfuncs.unnormalized_log_prob(yk)
    lqy = model.q.log_prob(y)
    print(tf.stack([lpx, lpy, lqx, lqy], axis=0).numpy())
    prob = tf.exp((lpy - lpx)/nc**3 + lqx - lqy)
    #prob = tf.exp((lpy - lpx + lqx - lqy))

    accept = tf.where(tf.random.uniform([1]) <= tf.minimum(1., prob))
    reject = tf.where(tf.random.uniform([1]) > tf.minimum(1., prob))
    z = tf.scatter_nd(accept, tf.gather_nd(yk, accept), yk.shape) + \
         tf.scatter_nd(reject, tf.gather_nd(xk, reject), yk.shape)

    lpz = tf.where(tf.random.uniform([1]) <= tf.minimum(1., prob), lpy, lpx)
    acc = tf.where(tf.random.uniform([1]) <= tf.minimum(1., prob), tf.ones(lpx.shape)*11, tf.ones(lpx.shape)*10.)
    return z, lpz, prob, acc, tf.constant(1.)

def mcmc_step(q, stepsize):
    #stepsize = np.random.uniform(0.01, 0.02, 1)
    lpsteps = np.random.randint(lpsteps1, lpsteps2, 1)[0]
    print("mcmc : ", q.shape)
    q, _, acc, prob, _ = hmckernel.hmc_step(q.numpy(), lpsteps, stepsize)
    lpq = tf.constant(prob[2]*-1.)
    prob = np.exp(prob[0] - prob[1])
    q = tf.constant(q, dtype=tf.float32)
    return q, lpq, tf.constant(prob), tf.constant(acc*1.), tf.constant(0.)

#@tf.function
def fvi(model, q, logpq, i, stepsize):

    elbos = []
    start = time.time()
    if i < mcmciter:
        print("MCMC only")
        q, logpq, prob, acc, jump =  mcmc_step(q, stepsize)
    else:
        q, logpq, prob, acc, jump = tf.cond(tf.random.uniform([1]) > probjump, \
                                      lambda : mcmc_step(q, stepsize), \
                            lambda : flow_step(model, q, logpq))
        if jump == 1:
            for _ in range(recalib):
                print("Jump, recaliberate")
                q, logpq, prob2, acc2, jump2 =  mcmc_step(q, stepsize)

    print("prob ", prob.numpy(), "to accept/reject with ", acc.numpy(), " in time %0.3f"%(time.time() - start))
    return q, logpq, acc, jump

@tf.function
def fvi_grads(model, samples, opt, samples2, regwt):
    samples = evolve.zk_to_z(samples)
    with tf.GradientTape(persistent=True) as tape:
        tape.watch(model.trainable_variables)
        #q, logpq, acc = fvi_step(model, q, logpq)
        neglogq = - tf.reduce_mean(model.q.log_prob(samples))
        if regwt0 ==  0:
            loss = neglogq*1.
        else:
            logq2 = - tf.reduce_mean(model.q.log_prob(evolve.zk_to_z(samples2)))
            samplevi = model.sample(samples.shape[0])*1.
            logqvi = - tf.reduce_mean(model.q.log_prob(samplevi))
            reg = regwt * tf.reduce_mean((logq2 - logqvi)**2.)
            loss = neglogq + reg

    gradients = tape.gradient(loss, model.trainable_variables)
    return neglogq, gradients

def fvi_train(model, samples, opt, batch2, regwt):
    neglogq, gradients = fvi_grads(model, samples, opt, batch2, regwt)
    opt.apply_gradients(zip(gradients, model.trainable_variables))
    return neglogq

def optimizetrenf(model, data, opt,  Rsm, niter=100, callback=callback, citer=50, nsamples=1, saveiter=20, sample=None, batch_size=1, stepsize=0.01, samples=[]):

    pss, psslin, pssx = [], [], []
    if sample is None:
        sample = tf.random.normal([nsamples, nc, nc, nc])
    logpsample = dmfuncs.unnormalized_log_prob(sample)

    accs, losses = [], []
    samples.append(sample)
    idxtrain = list(np.arange(len(samples)))
    #
    start = time.time()
    #dry run
    sample, logpsample, acc, jump = fvi(model, sample, logpsample, 0, stepsize)
    print(sample.shape, samples[0].shape, samples[-1].shape)
    accs.append(acc)
    #samples.append(sample)
    #idxtrain = idxtrain + [len(samples)-1]
    trainsize = len(samples)
    print("Training size of %d is maintained"%trainsize)
    for epoch in range(niter+1):
        print(epoch)
        sample, logpsample, acc, jump = fvi(model, sample, logpsample, epoch, stepsize)
        accs.append(acc)
        #save power spectrum
        if sample.shape[-1] == 2:
            ssz = evolve.zk_to_z(sample).numpy()
        else: ssz = tf.identity(sample)
        sslin = evolve.z_to_lin(tf.constant(ssz)).numpy()
        pss.append([tools.power(j+1, boxsize=evolve.bs)[1] for j in ssz])
        psslin.append([tools.power(j+1, boxsize=evolve.bs)[1] for j in sslin])
        pssx.append([tools.power(j+1, ic[0] + 1, boxsize=evolve.bs)[1] for j in sslin])
        samples.append(sample)
        idxtrain = idxtrain + [len(samples)-1]
        print(max(idxtrain), len(idxtrain), len(samples))
        if (epoch > 0) & (epoch %100 == 0):
            ntrain = 500
            print("Train %d times to recaliberate"%ntrain)
        else: ntrain = args.ntrain
        for itrain in range(ntrain):
            if jump == 1:
                print("Jump, no train")
                break
            if itrain == 0: idxsample = idxtrain[-batch_size:]
            else: idxsample = np.random.choice(idxtrain, batch_size, replace=False)
            #$# idxsample = np.random.choice(idx, batch_size, p=idx/idx.sum(), replace=False)
            idlast = min(-1 * np.random.choice(10), -1)
            idxsample = [idlast] + list(idxsample)
            print(itrain, ", idx :", idxsample)
            batch = tf.concat([samples[i] for i in idxsample], axis=0)
            batch2, regwt = None, tf.constant(regwt0)
            if regwt0 != 0.:
                if itrain == 0: idxsample = idxtrain[-batch_size:]
                else: idxsample = np.random.choice(idxtrain, batch_size, replace=False)
                #idxsample = idxtrain[np.random.choice(np.arange(trainsize), batch_size, replace=False)]
                idxsample = list(idxsample)
                print("idx2 :", idxsample)
                batch2 = tf.concat([samples[i] for i in idxsample], axis=0)
            loss = fvi_train(model, batch, opt, batch2, regwt)
            losses.append(loss/nc**3)
        print("Logq at epoch %d is: "%epoch, loss.numpy())
        #iddel = np.random.randint(trainsize)
        #print("Remove index from train : %d "%idxtrain[iddel])
        #del idxtrain[iddel]

        if epoch%citer == 0:
            print("pss shape ", pss[-1][0].shape, len(pss[-1]))
            print("Time taken for %d iterations : "%citer, time.time() - start)
            print("Accept counts : ", np.unique(accs, return_counts=True))
            #
            fig = callback_fvi(model, ic, bs, losses)
            plt.savefig(fpath + '/figs/iter%05d'%epoch)
            plt.close()

            #
            fig = callback_sampling([evolve.zk_to_lin(i) for i in samples[-2:]], ic, bs)
            plt.savefig(fpath + '/figs/samples%05d'%epoch)
            plt.close()
            np.save(fpath + 'loss', np.array(losses))
            np.save(fpath + 'samples', np.array(samples))
            np.save(fpath + '/psz', pss)
            np.save(fpath + '/pslin', psslin) #
            np.save(fpath + '/pslinx', pssx) #
            np.save(fpath + 'samplesthin%d'%thinning, np.array(samples)[::thinning])
            np.save(fpath + 'accepts', np.array(accs))
            start = time.time()
        if epoch%saveiter == 0:
            model.save_weights(fpath + '/weights/iter%04d'%(epoch//saveiter))
            #np.save(fpath + '/opt/iter%04d'%(epoch//saveiter), opt.get_weights(), allow_pickle=True)
            #print('Weights saved')

    return losses

#$####################################################################################
#$################################################################
###################VI
print("\nStart VI\n")
#vitrenf = trenfmodel.VItrenf(nc, nlayers=args.nlayers, evolve=evolve, nbins=32, mode=args.mode)
vitrenf = trenfmodel.VItrenf(nc, nlayers=args.nlayers, evolve=evolve, nbins=args.nbins, nknots=args.nknots, mode=args.mode, \
                             linknots=bool(args.linknots), fitnoise=bool(args.fitnoise), fitscale=bool(args.fitscale), \
                             fitmean=bool(args.fitmean), meanfield=bool(args.meanfield))
for i in vitrenf.variables:
    print(i.name, i.shape)

lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
    5e-3,
    decay_steps=200,
    decay_rate=0.95,
    staircase=True)

#Train FLOW using previos samples assuming we have them
print("Load old samples")
samplepath = '/mnt/ceph/users/cmodi/galference/dm_hmc/L1000_N0128_ZA-kwts4-fourier-corr/'
#samplepath = '/mnt/ceph/users/cmodi/galference/dm_hmc/L0200_N0064_ZA-fourier-minlininit2/'
allsamples = []
perchain = 201
for i in range(nchains):
    f = np.load(samplepath + '/samples%d-00.npy'%i)[:perchain, 0]
    allsamples.append(f)
allsamples = np.stack(allsamples, axis=1).astype(np.float32)
print(allsamples.shape)
stepsize = np.array([np.load(samplepath + '/stepsizes%d-00.npy'%i) for i in range(nchains)]).flatten()
print("stepsize : ", stepsize)

losses = []

opt = tf.keras.optimizers.Adam(learning_rate= lr_schedule)
#opt = tf.keras.optimizers.RMSprop(learning_rate= 1e-3)
saveiter = 100

print("Train on old samples")
for i in range(prentrain+1):

    for itrain in range(min(ntrain, int(ntrain*(2*i/prentrain)))):

        #idx = np.random.choice(i + 1, size=nsamples+1)
        idx = np.random.choice(allsamples.shape[0], size=nsamples+1)
        #idx = [-1] + list(idx)
        batch = tf.concat([allsamples[i] for i in idx], axis=0)
        #batch = tf.concat(allsamples[idx], axis=0)
        batch2, regwt = None, None
        if regwt0 !=0 :
            batch2 = tf.concat(allsamples[np.random.choice(allsamples.shape[0], nsamples)], axis=0)
            regwt = tf.constant(args.regwt0*i/prentrain)
        losses.append(fvi_train(vitrenf, batch, opt, batch2, regwt))
        #if itrain == 0: print(idx, batch.shape)
    if i%10 == 0:
        print("Iter : ", i)
        fig = callback_fvi(vitrenf, ic, bs, losses)
        plt.savefig(fpath + 'figs/trainr%05d'%i)
        plt.close()

#$#        lqs = []
#$#        for j in range(allsamples.shape[0]//10):
#$#            batch = tf.concat(evolve.zk_to_z(allsamples[j:j+2]), axis=0)
#$#            lq = vitrenf.q.log_prob(batch).numpy()
#$#            lqs = lqs + list(lq.flatten())
#$#        plt.plot(lqs, '.', label='train samples', lw=2)
#$#
#$#        lqs = []
#$#        for j in range(allsamples.shape[0]//10):
#$#            batch = vitrenf.sample(nchains)
#$#            lq = vitrenf.q.log_prob(batch).numpy()
#$#            lqs = lqs + list(lq.flatten())
#$#        plt.plot(lqs, '.', label='flow samples', lw=2)
#$#        plt.legend()
#$#        plt.savefig(fpath + '/figs/lgoqtrain%05d'%i)
#$#        plt.close()
#$#
    if (i > 0) & (i%saveiter == 0):
        vitrenf.save_weights(fpath + '/weights/preiter%04d'%(i//saveiter))

plt.imshow(vitrenf.sample(1)[0].numpy().sum(axis=0))
plt.colorbar()
plt.savefig(fpath + 'initsample')
plt.colorbar()
plt.close()

sample = tf.constant(allsamples[-1])
print(sample.shape)
R = 0.

opt2 = tf.keras.optimizers.Adam(learning_rate= 1e-3)
batch = tf.concat([allsamples[i] for i in idx], axis=0)
_ = fvi_train(vitrenf, batch, opt2, batch2, regwt)
opt2.set_weights(opt.get_weights())

losses = optimizetrenf(vitrenf, tf.constant(data), opt2, tf.constant(R), 5001, callback, citer=10, sample=sample, stepsize=stepsize, samples=list(allsamples[::thinning]))

#$#            #
#$#            lqs = []
#$#            for j in range(10):
#$#                idxsample = np.random.choice(idxtrain, batch_size, replace=False)
#$#                idxsample = list(idxsample)
#$#                try:
#$#                    print(idxsample)
#$#                    batch = tf.concat([samples[i:i+2] for i in list(idxsample)], axis=0)
#$#                    batch = evolve.zk_to_z(batch)
#$#                    lq = model.q.log_prob(batch).numpy()
#$#                    plt.plot(lq.flatten(), 'C0.')
#$#                    lqs = lqs + list(lq.flatten())
#$#                    #plt.plot(lqs, '.', label='train samples', lw=2)
#$#                except Exception as e:
#$#                    print(e)
#$#            lqs = []
#$#            for j in range(10):
#$#                try:
#$#                    batch = model.sample(nchains)
#$#                    lq = model.q.log_prob(batch).numpy()
#$#                    lqs = lqs + list(lq.flatten())
#$#                    plt.plot(lq.flatten(), 'C1.', lw=2)
#$#                except Exception as e:
#$#                    print(e)
#$#            plt.savefig(fpath + '/figs/lgoqiter%05d'%epoch)
#$#            plt.close()
#$#