from __future__ import division
import torch

# Beam decoding with KL divergence
import numpy as np
from scipy.stats import entropy

class Beam(object):
    """
    Class for managing the internals of the beam search process.

    Takes care of beams, back pointers, and scores.

    Args:
        size (int): beam size
        pad, bos, eos (int): indices of padding, beginning, and ending.
        vocab (vocab): vocab of the target
        syntax_topics_model (Syntax and Topics module): This is an object of the class which will have the topics
                                                        and classes word probabilities
        source (list): list of source indices which will be the source sentence
        targets (list of list): list of taget sentences each of which is a list of indices of the full hypothesis generated till now
        n_best (int): nbest size to use
        cuda (bool): use gpu
        global_scorer (:obj:`GlobalScorer`)
    """
    def __init__(self, size, pad, bos, eos,
                 vocab, syntax_topics_model,
                 source,
                 n_best=1, cuda=False,
                 global_scorer=None,
                 min_length=0):

        self.size = size
        self.tt = torch.cuda if cuda else torch

        # The score for each translation on the beam.
        self.scores = self.tt.FloatTensor(size).zero_()
        self.all_scores = []

        # The backpointers at each time-step.
        self.prev_ks = []

        # The outputs at each time-step.
        self.next_ys = [self.tt.LongTensor(size)
                        .fill_(pad)]
        self.next_ys[0][0] = bos
        ##NOTE: speical marker which tells which class is the previous word from
        self.next_ys_topic_prior_sum = []        # Store the sum of the topic prior for the entire hypothesis
        self.next_ys_class_prior_sum = []        # Store the sum of the class prior for the entire hypothesis
        self.TOPIC_FLAG = 0
        self.CLASS_FLAG = 1

        # Has EOS topped the beam yet.
        self._eos = eos
        self.eos_top = False

        # The attentions (matrix) for each time.
        self.attn = []

        # Time and k pair for finished.
        self.finished = []
        self.n_best = n_best

        # Information for global scoring.
        self.global_scorer = global_scorer
        self.global_state = {}

        # Minimum prediction length
        self.min_length = min_length

        ##NOTE: Custom code
        self.vocab = vocab
        self.finished_marker = [-1]*size
        self.syntax_topics_model = syntax_topics_model
        self.source = [self.vocab.itos[word_id] for word_id in source]
        # Compute the topic prior probability for the source sentence
        self.source_topic_prior = np.zeros(self.syntax_topics_model.num_topics+1, dtype=np.float)
        print self.source
        self.src_topic_word_count = 0
        for word in self.source:
            word_topic_prior = self.syntax_topics_model.get_topic_prior_for_word(word)
            if word_topic_prior[self.syntax_topics_model.num_topics] != 1.0:
                print word
                self.src_topic_word_count += 1
                self.source_topic_prior += word_topic_prior
        self.source_class_prior = np.zeros(self.syntax_topics_model.num_classes, dtype=np.float)
        for word in self.source:
            self.source_class_prior += self.syntax_topics_model.get_class_prior_for_word(word)

        self.source_topic_prior /= len(self.source)
        self.source_class_prior /= len(self.source)
        self.beta = 0.0         # Additive smoothing
        self.source_topic_prior += self.beta
        self.source_class_prior += self.beta
        print(self.source_topic_prior)
        print(self.source_class_prior)
        self.L = 50         # Number of words to be chosen for the similarity consideration
        # temp = np.zeros_like(self.source_topic_prior) + self.beta
        # least_KL = entropy(self.source_topic_prior, temp)
        self.alpha = 1.5                # Multiplicative factor for topic KL divergence
        self.gamma = 1.5                 # Multiplicative factor for class KL divergence
        # self.alpha = 8/least_KL        # multiplicative factor for the KL divergence
        # print vocab.itos[0]
        # print vocab.itos[bos]
        # print vocab.itos[self._eos]
        # print vocab.__dict__.keys()
        # print vocab.unk_init
        # print type(vocab)
        self.topic_KL_flag = False
        self.class_KL_flag = True

    def print_hyp(self, hyp, class_or_topic = None):
        for i, word_id in enumerate(hyp):
            if class_or_topic:
                print "{}_{} ".format(self.vocab.itos[word_id], class_or_topic[i]),
            else:
                print "{} ".format(self.vocab.itos[word_id]),

    ##NOTE: Custom function for debugging
    def print_all_words_in_beam(self):
        # Uses the vocab and prints the list of next_ys
        for i in range(self.size):
            timestep = len(self.next_ys)
            # if self.finished_marker[i] != -1:
            #     timestep = self.finished_marker[i]
            if timestep > 1:
                hyp, _ = self.get_hyp(timestep, i)
                # print hyp
                # print type(hyp)
                self.print_hyp(hyp)
                print "$$$"
            # print ""

    def print_received_targets(self):
        for i in range(len(self.targets)):
            self.print_hyp(self.targets[i])
            print ""
        print "############### Received Targets ############"

    def print_the_top_choices(self, best_choices):
        w, h = best_choices.size()
        for i in range(w):
            for j in range(h):
                print "{} ".format(self.vocab.itos[best_choices[i,j]]),
            print ""

    def get_current_state(self):
        "Get the outputs for the current timestep."
        return self.next_ys[-1]

    def get_current_origin(self):
        "Get the backpointers for the current timestep."
        return self.prev_ks[-1]

    def advance(self, word_probs, attn_out):
        print "Advancing beam"
        """
        Given prob over words for every last beam `wordLk` and attention
        `attn_out`: Compute and update the beam search.

        Parameters:

        * `word_probs`- probs of advancing from the last step (K x words)
        * `attn_out`- attention at the last step

        Returns: True if beam search is complete.
        """
        num_words = word_probs.size(1)

        # force the output to be longer than self.min_length
        cur_len = len(self.next_ys)
        if cur_len < self.min_length:
            for k in range(len(word_probs)):
                word_probs[k][self._eos] = -1e20

        # Sum the previous scores.
        if len(self.prev_ks) > 0:
            beam_scores = word_probs + \
                self.scores.unsqueeze(1).expand_as(word_probs)

            # Don't let EOS have children.
            for i in range(self.next_ys[-1].size(0)):
                if self.next_ys[-1][i] == self._eos:
                    beam_scores[i] = -1e20
        else:
            beam_scores = word_probs[0]

        #TODO: 1) find the last word for each beam
        #       2) For each possible hypothesis find the Topic modeling probability
        #       3) Weighted add the syntax and topic scores to beam_scores and just calculate the next_ys and backpointers as normal

        if len(self.prev_ks) > 0:
            per_beam_words = self.next_ys[-1]
            if self.topic_KL_flag:
                per_beam_hyp_topic_prior_sum = self.next_ys_topic_prior_sum[-1]
            if self.class_KL_flag:
                per_beam_hyp_class_prior_sum = self.next_ys_class_prior_sum[-1]

            if self.topic_KL_flag:
                topic_KL_divergence_scores = self.tt.zeros_like(beam_scores)
            if self.class_KL_flag:
                class_KL_divergence_scores = self.tt.zeros_like(beam_scores)
            for i in range(self.size):
                if self.topic_KL_flag:
                    hyp_topic_prior = per_beam_hyp_topic_prior_sum[i]
                if self.class_KL_flag:
                    hyp_class_prior = per_beam_hyp_class_prior_sum[i]

                len_hyp = len(self.next_ys)             # Includes current word because we want to skip the count added by the start word <bos>
                for j in range(num_words):
                    word = self.vocab.itos[j]
                    #KL divergence for Topic Priors
                    if self.topic_KL_flag:
                        word_topic_prior = self.syntax_topics_model.get_topic_prior_for_word(word)
                        hyp_topic_prior_sum = (hyp_topic_prior + word_topic_prior) + self.beta
                        # topic_KL_divergence_scores[i][j] = entropy(self.source_topic_prior, hyp_topic_prior_sum)
                        topic_KL_divergence_scores[i][j] = self.syntax_topics_model.KL_divergence(self.source_topic_prior, hyp_topic_prior_sum)
                        if topic_KL_divergence_scores[i][j] == float('Inf'):
                            print word
                            print topic_KL_divergence_scores[i][j]
                            print self.source_topic_prior
                            print hyp_topic_prior_sum
                    #KL divergence for Class Priors
                    if self.class_KL_flag:
                        if "<unk>" in word:
                            class_KL_divergence_scores[i][j] = 1000000.0
                            continue
                        word_class_prior = self.syntax_topics_model.get_class_prior_for_word(word)
                        hyp_class_prior_sum = (hyp_class_prior + word_class_prior) + self.beta
                        # class_KL_divergence_scores[i][j] = entropy(self.source_class_prior, hyp_class_prior_sum)
                        class_KL_divergence_scores[i][j] = self.syntax_topics_model.KL_divergence(self.source_class_prior, hyp_class_prior_sum)
                        if class_KL_divergence_scores[i][j] == float('Inf'):
                            print word
                            print class_KL_divergence_scores[i][j]
                            print self.source_class_prior
                            print hyp_class_prior_sum

            #TODO: Convert the zeros to max KL divergence
            if self.topic_KL_flag:
                max_topic_KL = topic_KL_divergence_scores.mean()
                for i in range(self.size):
                    for j in range(num_words):
                        word = self.vocab.itos[j]
                        if "<unk>" in word:
                            topic_KL_divergence_scores[i][j] = 1000000.0
                            continue
                        if topic_KL_divergence_scores[i][j] == 0.0:
                            topic_KL_divergence_scores[i][j] = max_topic_KL
                print "########\n", max_topic_KL, "\n#########"
                    # Manually discourage unk by setting large negative syntax_topic_probability
                    # if "<unk>" in word:
                    #     syntax_topic_score, best_class_or_topic, class_or_topic = -1000000.0, -1, "C"
                    # else:
                    #     syntax_topic_score, best_class_or_topic, class_or_topic = self.syntax_topics_model.get_log_prob(word, 0 if per_beam_words_class_or_topic[i] == self.TOPIC_FLAG else per_beam_words_class_or_topic_number[i])
            overall_scores = beam_scores - ((self.alpha * topic_KL_divergence_scores) if self.topic_KL_flag else 0.0) - ((self.gamma * class_KL_divergence_scores) if self.class_KL_flag else 0.0)
            # print "Overall Score"
            # print overall_scores
            # print overall_scores.size()

            size = int(overall_scores.size(1))
            # print "Size of the overall_scores = ", size
            flat_beam_scores = overall_scores.view(-1)
            best_scores, best_scores_id = flat_beam_scores.topk(self.size, 0,
                                                                True, True)

            # We will debug the individual scores of the best candidates
            word_prob_best_scores = self.tt.zeros_like(best_scores)
            for i in range(self.size):
                word_prob_best_scores[i] = beam_scores[int(best_scores_id[i] / size)][best_scores_id[i] - int(best_scores_id[i] / size) * size]
            if self.topic_KL_flag:
                topic_KL_divergence_best_scores = self.tt.zeros_like(best_scores)
                for i in range(self.size):
                    topic_KL_divergence_best_scores[i] = topic_KL_divergence_scores[int(best_scores_id[i] / size)][best_scores_id[i] - int(best_scores_id[i] / size) * size]
            if self.class_KL_flag:
                class_KL_divergence_best_scores = self.tt.zeros_like(best_scores)
                for i in range(self.size):
                    class_KL_divergence_best_scores[i] = class_KL_divergence_scores[int(best_scores_id[i] / size)][best_scores_id[i] - int(best_scores_id[i] / size) * size]
            print best_scores
            print word_prob_best_scores
            # print topic_KL_divergence_best_scores
            # print class_KL_divergence_best_scores
            if self.topic_KL_flag:
                print self.alpha
                print self.alpha * topic_KL_divergence_best_scores
            if self.class_KL_flag:
                print self.gamma
                print self.gamma * class_KL_divergence_best_scores

            if self.topic_KL_flag:
                KL_size = int(topic_KL_divergence_scores.size(1))
                # print "Size of the overall_scores = ", size
                flat_beam_scores = topic_KL_divergence_scores.view(-1)
                debug_size = 25
                best_topic_KL_scores, best_topic_KL_scores_id = flat_beam_scores.topk(debug_size, 0,
                                                                False, True)
                print "Best words from topic KL"
                for i in range(debug_size):
                    # print best_topic_KL_scores_id[i], best_topic_KL_scores_id[i] - int(best_topic_KL_scores_id[i] / KL_size) * KL_size
                    word = self.vocab.itos[best_topic_KL_scores_id[i] - int(best_topic_KL_scores_id[i] / KL_size) * KL_size]
                    word_topic_prior = self.syntax_topics_model.get_topic_prior_for_word(word)
                    # print word, word_topic_prior, entropy(self.source_topic_prior, word_topic_prior + self.beta), \
                    print word, \
                                overall_scores[int(best_topic_KL_scores_id[i] / KL_size)][best_topic_KL_scores_id[i] - int(best_topic_KL_scores_id[i] / KL_size) * KL_size], \
                                best_topic_KL_scores[i], \
                                topic_KL_divergence_scores[int(best_topic_KL_scores_id[i] / KL_size)][best_topic_KL_scores_id[i] - int(best_topic_KL_scores_id[i] / KL_size) * KL_size]
                                # class_KL_divergence_scores[int(best_topic_KL_scores_id[i] / KL_size)][best_topic_KL_scores_id[i] - int(best_topic_KL_scores_id[i] / KL_size) * KL_size], \
                # print word_prob_best_scores + self.alpha * syntax_topic_best_scores

            prev_k = best_scores_id / num_words
            # Update next_ys_topic_prior_sum for all beams
            if self.topic_KL_flag:
                best_hyp_topic_prior_sum = list()
                for i in range(self.size):
                    word = self.vocab.itos[best_scores_id[i] - int(best_scores_id[i] / size) * size]
                    word_topic_prior = self.syntax_topics_model.get_topic_prior_for_word(word)
                    if word_topic_prior[self.syntax_topics_model.num_topics] != 1.0:
                        best_hyp_topic_prior_sum.append(self.next_ys_topic_prior_sum[-1][int(best_scores_id[i] / size)] + word_topic_prior)
                    else:
                        # Add the word_topic_prior only if its a topic word
                        best_hyp_topic_prior_sum.append(self.next_ys_topic_prior_sum[-1][int(best_scores_id[i] / size)])
                self.next_ys_topic_prior_sum.append(best_hyp_topic_prior_sum)
            # Update next_ys_class_prior_sum for all beams
            if self.class_KL_flag:
                best_hyp_class_prior_sum = list()
                for i in range(self.size):
                    word = self.vocab.itos[best_scores_id[i] - int(best_scores_id[i] / size) * size]
                    word_class_prior = self.syntax_topics_model.get_class_prior_for_word(word)
                    best_hyp_class_prior_sum.append(self.next_ys_class_prior_sum[-1][int(best_scores_id[i] / size)] + word_class_prior)
                self.next_ys_class_prior_sum.append(best_hyp_class_prior_sum)
            self.prev_ks.append(prev_k)
            self.next_ys.append((best_scores_id - prev_k * num_words))
            self.attn.append(attn_out.index_select(0, prev_k))
            # exit()

            self.print_all_words_in_beam()
            print "############## After new words chosen ###########"
        else:
            # beam_scores is only V dimensional vector
            # Thus for every word add the KL divergence score to the beam score
            if self.topic_KL_flag:
                topic_KL_divergence_scores = self.tt.zeros_like(beam_scores)
            if self.class_KL_flag:
                class_KL_divergence_scores = self.tt.zeros_like(beam_scores)
            for i in range(num_words):
                word = self.vocab.itos[i]
                if "<unk>" in word:
                    if self.topic_KL_flag:
                        topic_KL_divergence_scores[i] = 1000000.0
                    if self.class_KL_flag:
                        class_KL_divergence_scores[i] = 1000000.0
                    beam_scores[i] = -1000000.0
                    continue
                # KL for Topic priors of the first word
                if self.topic_KL_flag:
                    word_topic_prior = self.syntax_topics_model.get_topic_prior_for_word(word)
                    word_topic_prior += self.beta
                    # topic_KL_divergence_scores[i] = entropy(self.source_topic_prior, word_topic_prior)
                    topic_KL_divergence_scores[i] = self.syntax_topics_model.KL_divergence(self.source_topic_prior, word_topic_prior)
                # KL for Class priors of the first word
                if self.class_KL_flag:
                    word_class_prior = self.syntax_topics_model.get_class_prior_for_word(word)
                    word_class_prior += self.beta
                    # class_KL_divergence_scores[i] = entropy(self.source_class_prior, word_class_prior)
                    class_KL_divergence_scores[i] = self.syntax_topics_model.KL_divergence(self.source_class_prior, word_class_prior)

            overall_scores = beam_scores - ((self.alpha * topic_KL_divergence_scores) if self.topic_KL_flag else 0.0) -  ((self.gamma * class_KL_divergence_scores) if self.class_KL_flag else 0.0)
            # flat_beam_scores = overall_scores.view(-1)
            flat_beam_scores = beam_scores.view(-1)             # For the first iteration use only the word probabilities
            best_scores, best_scores_id = flat_beam_scores.topk(self.size, 0,
                                                                True, True)

            self.all_scores.append(self.scores)
            # self.scores = best_scores           # will store the word_prob + the KL divergence for hypothesis
            self.scores = best_scores           # will store the word_prob log prob for hypothesis

            # best_scores_id is flattened, beam * word array, so calculate which
            # word and beam each score came from
            size = int(overall_scores.size(0))
            prev_k = best_scores_id / num_words
            # Update next_ys_topic_prior_sum for all beams
            if self.topic_KL_flag:
                best_hyp_topic_prior_sum = list()
                for i in range(self.size):
                    word = self.vocab.itos[best_scores_id[i]]
                    word_topic_prior = self.syntax_topics_model.get_topic_prior_for_word(word)
                    if word_topic_prior[self.syntax_topics_model.num_topics] != 1.0:
                        best_hyp_topic_prior_sum.append(word_topic_prior)
                    else:
                        # starting word is a syntax word. Therefore we will set the prior to zeros
                        best_hyp_topic_prior_sum.append(np.zeros((self.syntax_topics_model.num_topics+1), dtype=np.float))
                self.next_ys_topic_prior_sum.append(best_hyp_topic_prior_sum)

            # Update next_ys_class_prior_sum for all beams
            if self.class_KL_flag:
                best_hyp_class_prior_sum = list()
                for i in range(self.size):
                    word = self.vocab.itos[best_scores_id[i]]
                    word_class_prior = self.syntax_topics_model.get_class_prior_for_word(word)
                    best_hyp_class_prior_sum.append(word_class_prior)
                self.next_ys_class_prior_sum.append(best_hyp_class_prior_sum)

            self.prev_ks.append(prev_k)
            self.next_ys.append((best_scores_id - prev_k * num_words))
            self.attn.append(attn_out.index_select(0, prev_k))

        if self.global_scorer is not None:
            self.global_scorer.update_global_state(self)

        for i in range(self.next_ys[-1].size(0)):
            if self.next_ys[-1][i] == self._eos:
                s = self.scores[i]
                timestep = len(self.next_ys)
                if self.global_scorer is not None:
                    global_scores = self.global_scorer.score(self, self.scores)
                    s = global_scores[i]
                self.finished.append((s, len(self.next_ys) - 1, i))
                # TODO: Experimental!! Dividing the finished scores by their lenghts to be fair
                # self.finished.append((s/float(len(self.next_ys) - 1), len(self.next_ys) - 1, i))

                ##NOTE: Custom code
                if self.finished_marker[i] == -1:
                    # print "SET AND FORGET FOR ", i, "#$#$#$#$##$"
                    self.finished_marker[i] = len(self.next_ys) - 1

        # End condition is when top-of-beam is EOS and no global score.
        if self.next_ys[-1][0] == self._eos:
            # self.all_scores.append(self.scores)
            self.eos_top = True

        ##NOTE: Debugging
        # print word_probs, "$$"
        # print self.get_current_state(), "$$"
        # self.print_all_words_in_beam()
        # print "############## Beam Advance ###########"

    def done(self):
        return self.eos_top and len(self.finished) >= self.n_best

    def sort_finished(self, minimum=None):
        if minimum is not None:
            i = 0
            # Add from beam until we have minimum outputs.
            while len(self.finished) < minimum:
                s = self.scores[i]
                if self.global_scorer is not None:
                    global_scores = self.global_scorer.score(self, self.scores)
                    s = global_scores[i]
                self.finished.append((s, len(self.next_ys) - 1, i))

        # print self.finished
        # exit()

        self.finished.sort(key=lambda a: -a[0])
        scores = [sc for sc, _, _ in self.finished]
        ks = [(t, k) for _, t, k in self.finished]
        return scores, ks

    def get_hyp_with_class(self, timestep, k):
        """
        Walk back to construct the full hypothesis while also storing the class/topic number
        """
        hyp, class_or_topic = [], []
        # print len(self.next_ys), len(self.next_class_or_topics), len(self.next_class_or_topic_numbers)
        for j in range(len(self.prev_ks[:timestep]) - 1, -1, -1):
            hyp.append(self.next_ys[j+1][k])
            class_or_topic.append("{}{}".format("C" if self.next_class_or_topics[j][k] == 1 else "T", self.next_class_or_topic_numbers[j][k]))
            k = self.prev_ks[j][k]
        class_or_topic.reverse()
        return hyp[::-1], class_or_topic

    def get_hyp(self, timestep, k):
        """
        Walk back to construct the full hypothesis.
        """
        hyp, attn = [], []
        for j in range(len(self.prev_ks[:timestep]) - 1, -1, -1):
            hyp.append(self.next_ys[j+1][k])
            attn.append(self.attn[j][k])
            k = self.prev_ks[j][k]
        return hyp[::-1], torch.stack(attn[::-1])

class GNMTGlobalScorer(object):
    """
    NMT re-ranking score from
    "Google's Neural Machine Translation System" :cite:`wu2016google`

    Args:
       alpha (float): length parameter
       beta (float):  coverage parameter
    """
    def __init__(self, alpha, beta):
        self.alpha = alpha
        self.beta = beta

    def score(self, beam, logprobs):
        "Additional term add to log probability"
        cov = beam.global_state["coverage"]
        pen = self.beta * torch.min(cov, cov.clone().fill_(1.0)).log().sum(1)
        l_term = (((5 + len(beam.next_ys)) ** self.alpha) /
                  ((5 + 1) ** self.alpha))
        return (logprobs / l_term) + pen

    def update_global_state(self, beam):
        "Keeps the coverage vector as sum of attens"
        if len(beam.prev_ks) == 1:
            beam.global_state["coverage"] = beam.attn[-1]
        else:
            beam.global_state["coverage"] = beam.global_state["coverage"] \
                .index_select(0, beam.prev_ks[-1]).add(beam.attn[-1])


"""
Joseph Cook, Aarhus University, Jan 2022

This script contains functions that a) generate
SNICAR-predicted spectral albedo that approximate
field-measured spectra for a variety of weathering
crust configurations; b) quantify the albedo change
resulting from a range of WC development scenarios

includes:

find_best_params()
    the forward modelling script that retrieves the snicar params that generate the
    best-matching curve to a given field spectrum

call_snicar()
    the function used to do multiple snicar runs with params provided as a named tuple

match_field_spectra()
    the function for plotting simulated and measured field spectra in a multipanel fig

isolate_biological_effect()
    the function for calculating the albedo reduction due to bio vs phys processes by spectral differencing

build_LUT()
    function for constructing lookup table to be used in the inverse model

inverse_model()
    function finds the best matching entry in the LUTs for each field spectrum and returns the params
    used to generate it

"""
import collections as c
import sys

sys.path.append("./src")
import dask
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import statsmodels.api as sm
import xarray as xr
from call_snicar import call_snicar

def match_field_spectra(
    FIELD_DATA_FNAME,
    fnames,
    rho,
    rds,
    dz,
    alg,
    measured_cells,
    CIsites,
    LAsites,
    HAsites,
    APPLY_ARF,
    PLOT_ARF,
    ARF_CI,
    ARF_HA,
    SAVEPATH,
):

    """
    plot field against SNICAR spectra
    requires parameters to be known in advance and hard coded inside this function
    the relevant params can be generated using the find_best_params() func

    params:
    FIELD_DATA_FNAME = filename for spectral database

    The following params are parallel arrays - the order matters and
    must match the filenames! Pairs of values represent values for
    upper and lower layers in model. These are the values used to
    generate snicar spectrum to match field spectrum with name = fname[i]

    fnames = sample IDs for spectra to match model runs
    rho = pairs of density values [upper, lower]
    rds = pairs of r_eff values [upper, lower]
    dz = layer thickness [upper, lower] NB. upper is always 0.001
    alg = mass concentration of algae [upper, lower] NB. lower is always 0
    measured_cells = array of the actual measured cell concentration for each spectrum

    e.g.
    fnames= ['2016_WI_8','14_7_SB6','14_7_SB9','14_7_SB1','21_7_SB2','14_7_SB2', '22_7_SB3', 'RAIN']
    rho = [[550,550],[650,650],[800,800],[850,850],[750,750],[800,800],[800,800],[900,900]]
    rds = [[550,550],[650,650],[850,850],[850,850],[800,800],[800,800],[750,750],[900,900]]
    dz = [[0.001,0.3],[0.001,0.09],[0.001,0.03],[0.001,0.03],[0.001,0.02],[0.001,0.06],[0.001,0.05],[0.001,0.03]]
    alg = [[0,0],[0,0],[20000,0],[30000,0],[45000,0],[3000,0],[8000,0],[0,0]]

    returns:
    None, but saves figure to SAVEPATH

    """

    spectra = pd.read_csv(FIELD_DATA_FNAME)

    # reformat feld spectra to match snicar resolution
    spectra = spectra[::10]

    # gather spectra for each surface type
    CIspec = spectra[spectra.columns.intersection(CIsites)]
    HAspec = spectra[spectra.columns.intersection(HAsites)]
    LAspec = spectra[spectra.columns.intersection(LAsites)]
    RAINspec = spectra["RAIN2"]

    if PLOT_ARF:
        plt.plot(
            spectra.Wavelength[0:100], ARF_CI[0:100], marker="x", label="clean ice ARF"
        ),
        plt.plot(
            spectra.Wavelength[0:100],
            ARF_HA[0:100],
            marker="o",
            linestyle="dashed",
            label="algal ice ARF",
        )
        plt.ylabel("Anitostropic Reflectance Factor"), plt.xlabel("Wavelength (nm)")
        plt.legend(loc="best")
        plt.savefig(str(SAVEPATH + "ARF.jpg"), dpi=300)

    # define local function for calling snicar
    def simulate_albedo(rds, rho, dz, alg):

        params = c.namedtuple(
            "params",
            "rho_layers, grain_rds, layer_type, dz,\
                 mss_cnc_glacier_algae, c_factor_GA, solzen",
        )
        params.grain_rds = rds
        params.rho_layers = rho
        params.layer_type = [1, 1]
        params.dz = dz
        params.mss_cnc_glacier_algae = alg
        params.c_factor_GA = 20
        params.solzen = 40
        albedo, BBA = call_snicar(params)

        return albedo, BBA

    # set up output array
    # and call snicar with each set of params
    OutArray = np.zeros(shape=(len(fnames), 480))

    for i in np.arange(0, len(fnames), 1):

        albedo, BBA = simulate_albedo(rds[i], rho[i], dz[i], alg[i])
        if APPLY_ARF:
            if alg[i][0] > 5000:
                albedo[15:230] = albedo[15:230] * ARF_HA
            else:
                albedo[15:230] = albedo[15:230] * ARF_CI

        OutArray[i, :] = albedo

    # calculate mean absolute error for model vs measured spectrum
    error = []
    for i in np.arange(0, len(fnames), 1):
        error.append(
            abs(np.mean(spectra[fnames[i]].iloc[0:130] - (OutArray[i, 15:145])))
        )

    # plot figure
    fig, axes = plt.subplots(4, 2, figsize=(10, 10))

    if APPLY_ARF:
        ylabel = "Reflectance"
        PlotName = "FieldvsMeasuredReflectance.jpg"
    else:
        ylabel = "Albedo"
        PlotName = "FieldvsMeasuredAlbedo.jpg"

    axes[0, 0].plot(spectra.Wavelength, spectra["23_7_SB1"], label="field")
    axes[0, 0].plot(
        spectra.Wavelength, OutArray[0, 15:230], label="model", linestyle="--"
    )
    axes[0, 0].set_ylim(0, 1), axes[0, 0].set_xlim(350, 1800)
    axes[0, 0].text(
        1400,
        0.2,
        "r_eff: {}\nrho: {}\ndz: {}\nC-factor: 30\nerror: {:.3f}".format(
            rds[0][0], rho[0][0], dz[0][1], error[0]
        ),
    )
    axes[0, 0].text(400, 0.1, "{} \n{} cells/mL ".format(fnames[0], measured_cells[0]))
    axes[0, 0].set_ylabel(ylabel), axes[0, 0].set_xlabel("Wavelength (nm)")
    axes[0, 0].legend(loc="best")

    axes[0, 1].plot(spectra.Wavelength, spectra["14_7_SB6"])
    axes[0, 1].plot(spectra.Wavelength, OutArray[1, 15:230], linestyle="--")
    axes[0, 1].set_ylim(0, 1), axes[0, 1].set_xlim(350, 1800)
    axes[0, 1].text(
        1450,
        0.3,
        "r_eff: {}\nrho: {}\ndz: {}\nC-factor: 30\nerror: {:.3f}".format(
            rds[1][0], rho[1][0], dz[1][1], error[1]
        ),
    )
    axes[0, 1].text(400, 0.8, "{}: \n{} cells/mL".format(fnames[1], measured_cells[1]))
    axes[0, 1].set_ylabel(ylabel), axes[0, 1].set_xlabel("Wavelength (nm)")

    axes[1, 0].plot(spectra.Wavelength, spectra["14_7_SB9"])
    axes[1, 0].plot(spectra.Wavelength, OutArray[2, 15:230], linestyle="--")
    axes[1, 0].set_ylim(0, 1), axes[1, 0].set_xlim(350, 1800)
    axes[1, 0].text(
        1400,
        0.3,
        "r_eff: {}\nrho: {}\ndz: {}\nC-factor: 30\nerror: {:.3f}".format(
            rds[2][0], rho[2][0], dz[2][1], error[2]
        ),
    )
    axes[1, 0].text(400, 0.8, "{}: \n{} cells/mL".format(fnames[2], measured_cells[2]))
    axes[1, 0].set_ylabel(ylabel), axes[1, 0].set_xlabel("Wavelength (nm)")

    axes[1, 1].plot(spectra.Wavelength, spectra["14_7_SB1"])
    axes[1, 1].plot(spectra.Wavelength, OutArray[3, 15:230], linestyle="--")
    axes[1, 1].set_ylim(0, 1), axes[1, 1].set_xlim(350, 1800)
    axes[1, 1].text(
        1400,
        0.3,
        "r_eff: {}\nrho: {}\ndz: {}\nC-factor: 30\nerror: {:.3f}".format(
            rds[3][0], rho[3][0], dz[3][1], error[3]
        ),
    )
    axes[1, 1].text(400, 0.8, "{}: \n{} cells/mL".format(fnames[3], measured_cells[3]))
    axes[1, 1].set_ylabel(ylabel), axes[1, 1].set_xlabel("Wavelength (nm)")

    axes[2, 0].plot(spectra.Wavelength, spectra["22_7_SB5"])
    axes[2, 0].plot(spectra.Wavelength, OutArray[4, 15:230], linestyle="--")
    axes[2, 0].set_ylim(0, 1), axes[2, 0].set_xlim(350, 1800)
    axes[2, 0].text(
        1400,
        0.3,
        "r_eff: {}\nrho: {}\ndz: {}\nC-factor: 30\nerror: {:.3f}".format(
            rds[4][0], rho[4][0], dz[4][1], error[4]
        ),
    )
    axes[2, 0].text(400, 0.8, "{}: \n{} cells/mL".format(fnames[4], measured_cells[4]))
    axes[2, 0].set_ylabel(ylabel), axes[2, 0].set_xlabel("Wavelength (nm)")

    axes[2, 1].plot(spectra.Wavelength, spectra["14_7_SB2"])
    axes[2, 1].plot(spectra.Wavelength, OutArray[5, 15:230], linestyle="--")
    axes[2, 1].set_ylim(0, 1), axes[2, 1].set_xlim(350, 1800)
    axes[2, 1].text(
        1400,
        0.3,
        "r_eff: {}\nrho: {}\ndz: {}\nC-factor: 30\nerror: {:.3f}".format(
            rds[5][0], rho[5][0], dz[5][1], error[5]
        ),
    )
    axes[2, 1].text(400, 0.8, "{}: \n{} cells/mL".format(fnames[5], measured_cells[5]))
    axes[2, 1].set_ylabel(ylabel), axes[2, 1].set_xlabel("Wavelength (nm)")

    axes[3, 0].plot(spectra.Wavelength, spectra["22_7_SB3"])
    axes[3, 0].plot(spectra.Wavelength, OutArray[6, 15:230], linestyle="--")
    axes[3, 0].set_ylim(0, 1), axes[3, 0].set_xlim(350, 1800)
    axes[3, 0].text(
        1400,
        0.3,
        "r_eff: {}\nrho: {}\ndz: {}\nC-factor: 30\nerror: {:.3f}".format(
            rds[6][0], rho[6][0], dz[6][1], error[6]
        ),
    )
    axes[3, 0].text(400, 0.8, "{}: \n{} cells/mL".format(fnames[6], measured_cells[6]))
    axes[3, 0].set_ylabel(ylabel), axes[3, 0].set_xlabel("Wavelength (nm)")

    axes[3, 1].plot(spectra.Wavelength, RAINspec)
    axes[3, 1].plot(spectra.Wavelength, OutArray[7, 15:230], linestyle="--")
    axes[3, 1].set_ylim(0, 1), axes[3, 1].set_xlim(350, 1800)
    axes[3, 1].text(
        1400,
        0.3,
        "r_eff: {}\nrho: {}\ndz: {}\nC-factor: 30\nerror: {:.3f}".format(
            rds[7][0], rho[7][0], dz[7][1], error[7]
        ),
    )
    axes[3, 1].text(400, 0.8, "{}: \n{} cells/mL".format(fnames[1], measured_cells[1]))
    axes[3, 1].set_ylabel(ylabel), axes[3, 1].set_xlabel("Wavelength (nm)")

    fig.tight_layout()
    plt.savefig(str(SAVEPATH + PlotName), dpi=300)

    return True

def build_LUT(
    cfactor,
    solzen,
    dz,
    densities,
    radii,
    algae,
    wavelengths,
    save_LUT,
    APPLY_ARF,
    ARF_CI,
    ARF_HA,
    LUT_PATH,
):

    """
    generates LUTs used to invert BioSNICAR in RISA project

    params:
    ice_rds: fixed effective bubble radius for solid ice layers (default = 525)
    ice dens: fixed density for solid ice layers (default = 894)
    zeniths: range of solar zenith angles to loop over
    dz: thickness of each vertical layer
    densities: densities for top layer. Lower layers predicted by exponential model
    algae: mass mixing ratio of algae in top layer
    wavelengths: wavelength range, default is np.arange(0.2, 5, 0.01)
    save_LUT: Boolean to toggle saving to npy file
    SAVEPATH: directory to save LUT

    returns:
    WCthickLUT: for each index position in the spectraLUT, this holds the WC
                thickness in the corresponding index position
    SpectraLUT: ND array containing 480element spectrum for each
                dens/alg/zen combination

    return spectraLUT

    """
    LUT = []

    @dask.delayed
    def run_sims(cfactor, dens, rad, dz, alg, zen):

        params = c.namedtuple(
            "params",
            "rho_layers, grain_rds, layer_type, c_factor_GA, dz, mss_cnc_glacier_algae, solzen",
        )
        params.rho_layers = [916, dens]
        params.grain_rds = [rad, rad]  # set equal to density
        params.layer_type = [1, 1]
        params.dz = [0.001, dz]
        params.mss_cnc_glacier_algae = [alg, 0]
        params.solzen = zen
        params.c_factor_GA = cfactor

        albedo, BBA = call_snicar(params)

        return albedo

    for x in np.arange(0, len(cfactor), 1):
        for z in np.arange(0, len(solzen), 1):
            for i in np.arange(0, len(densities), 1):
                for j in np.arange(0, len(radii), 1):
                    for p in np.arange(0, len(dz), 1):
                        for q in np.arange(0, len(algae), 1):

                            albedo = run_sims(
                                cfactor[x],
                                densities[i],
                                radii[j],
                                dz[p],
                                algae[q],
                                solzen[z],
                            )

                            LUT.append(albedo)

    LUT = dask.compute(*LUT, num_workers=12)
    LUT = np.array(LUT).reshape(
        len(cfactor),
        len(solzen),
        len(densities),
        len(radii),
        len(dz),
        len(algae),
        len(wavelengths),
    )

    # move the ARF application to new loop because dask compute objets are immutable
    # i.e. modifications to albedo must be done post-compute
    if APPLY_ARF:
        for x in np.arange(0, len(cfactor), 1):
            for z in np.arange(0, len(solzen), 1):
                for i in np.arange(0, len(densities), 1):
                    for j in np.arange(0, len(radii), 1):
                        for p in np.arange(0, len(dz), 1):
                            for q in np.arange(0, len(algae), 1):

                                if algae[q] > 5000:

                                    LUT[x, z, i, j, p, q, 15:230] = (
                                        LUT[x, z, i, j, p, q, 15:230] * ARF_HA
                                    )

                                else:
                                    LUT[x, z, i, j, p, q, 15:230] = (
                                        LUT[x, z, i, j, p, q, 15:230] * ARF_CI
                                    )

    if save_LUT:
        np.save(str(LUT_PATH + "LUT.npy"), LUT)

    return LUT

def inverse_model(
    FIELD_DATA_FNAME, LUT_PATH, SITES, DZ, DENSITIES, RADII, ZENS, ALGAE, SAVEPATH
):
    """
    function takes arrays of vals used to build LUT and runs a 2 step inversion.
    First, each spectrum in the LUT is compared to each field spectrum
    The parameters giving the smallest mean error are selected
    Then those parameters are fixed and the process repeats, only varying algal concentration
    This gives the parameter set that best simulates each field measurement

    """
    # read in and reshape luts and field spectra
    field_data = pd.read_csv(FIELD_DATA_FNAME, index_col=None)
    field_spectra = field_data[::10]
    lut = np.load(str(LUT_PATH + "LUT.npy"))
    vis_start_idx = 15
    vis_end_idx = 55
    nir_start_idx = 55
    nir_end_idx = 230
    lut_nir = lut[:, :, :, :, :, nir_start_idx:nir_end_idx]
    lut_vis = lut[:, :, :, :, :, vis_start_idx:vis_end_idx]
    flat_nir_lut = lut_nir.reshape(
        len(ZENS) * len(DENSITIES) * len(RADII) * len(DZ) * len(ALGAE), 175
    )

    output = pd.DataFrame()
    retrieved_zen = []
    retrieved_density = []
    retrieved_radii = []
    retrieved_dz = []
    retrieved_algae = []
    names = []
    nir_errors = []
    vis_errors = []
    total_errors = []

    for (name, data) in field_spectra.iteritems():
        # filter to sites defined in config
        if name in SITES:
            names.append(name)

            # step 1: find params that match best in NIR
            error_array = np.sqrt(abs(flat_nir_lut**2 - np.array(data[40:]) ** 2))
            error_array = np.nan_to_num(
                error_array, nan=9999
            )  # protect agaiunst nans being interpreted as low error
            error_list = np.sum(error_array, axis=1)
            index = np.argmin(error_list)
            nir_errors.append(error_list[index])
            param_idx_phys = np.unravel_index(
                index, (len(ZENS), len(DENSITIES), len(RADII), len(DZ), len(ALGAE), 1)
            )

            # step 2: fix physical params and minimise error
            # in vis by varying algae only
            lut2 = lut_vis[
                param_idx_phys[0],
                param_idx_phys[1],
                param_idx_phys[2],
                param_idx_phys[3],
                :,
                :,
            ]

            np.sqrt(abs(flat_nir_lut**2 - np.array(data[40:]) ** 2))
            error_array = np.sqrt(
                abs(lut2**2 - np.array(data[vis_start_idx:vis_end_idx]) ** 2)
            )
            error_list = np.sum(error_array, axis=1)
            index = np.argmin(error_list)
            vis_errors.append(error_list[index])
            param_idx_alg = np.unravel_index(index, [1, 1, 1, 1, len(ALGAE), 1])

            # step 3: organize out data
            retrieved_zen.append(ZENS[param_idx_phys[0]])
            retrieved_density.append(DENSITIES[param_idx_phys[1]])
            retrieved_radii.append(RADII[param_idx_phys[2]])
            retrieved_dz.append(DZ[param_idx_phys[3]])
            retrieved_algae.append(ALGAE[param_idx_alg[4]])

            total_error = abs(
                lut[
                    param_idx_phys[0],
                    param_idx_phys[1],
                    param_idx_phys[2],
                    param_idx_phys[3],
                    param_idx_alg[4],
                    vis_start_idx:nir_end_idx,
                ]
                - data
            )
            total_errors.append(np.mean(total_error))

    output["fname"] = names
    output["solzen"] = retrieved_zen
    output["density"] = retrieved_density
    output["radii"] = retrieved_radii
    output["dz"] = retrieved_dz
    output["algae"] = retrieved_algae
    output["nir_error"] = nir_errors
    output["vis_error"] = vis_errors
    output["total_error"] = total_errors

    output.to_csv(str(SAVEPATH + "inverse_model_output.csv"))

    return output

def isolate_biological_effect(FIELD_DATA_FNAME, CIsites, LAsites, HAsites, SAVEPATH):

    """
    This function estimates the albedo reduction resulting from the ice physical changes
    versus the biological growth.

    Some nuance to the interpretation because the ce surface likely would not
    degrade to the same extent without the algal bloom.

    """

    # read in spectral database
    spectra = pd.read_csv(FIELD_DATA_FNAME)

    # reformat feld spectra to match snicar resolution
    spectra = spectra[::10]
    CIspec = spectra[spectra.columns.intersection(CIsites)]
    HAspec = spectra[spectra.columns.intersection(HAsites)]
    LAspec = spectra[spectra.columns.intersection(LAsites)]
    meanCI = CIspec.mean(axis=1)
    meanHA = HAspec.mean(axis=1)
    meanLA = LAspec.mean(axis=1)

    # define local function for calling snicar
    def simulate_albedo(rds, rho, dz, alg):
        params = c.namedtuple(
            "params",
            "rho_layers, grain_rds, layer_type, dz, mss_cnc_glacier_algae, solzen",
        )
        params.grain_rds = rds
        params.rho_layers = rho
        params.layer_type = [1, 1]
        params.dz = dz
        params.mss_cnc_glacier_algae = alg
        params.solzen = 53
        albedo, BBA = call_snicar(params)
        return albedo, BBA

    # call snicar to generat esimulated spectrum for LA and HA using predetermined params

    SNICARalbedoLA, BBA = simulate_albedo([800, 800], [800, 800], [0.001, 0.1], [0, 0])
    SNICARalbedoHA, BBA = simulate_albedo([900, 900], [800, 800], [0.001, 0.08], [0, 0])
    SNICARalbedoLA = SNICARalbedoLA[15:230]
    SNICARalbedoHA = SNICARalbedoHA[15:230]

    # plot figure
    x = np.arange(350, 2500, 10)
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))
    ax1.plot(x, meanCI, linestyle="--", alpha=0.4, label="Clean ice (mean)")
    ax1.plot(x, meanLA, linestyle="-.", alpha=0.4, label="Algal ice (mean)")
    ax1.plot(
        x, SNICARalbedoLA, linestyle="dotted", alpha=0.4, label="Clean ice (model)"
    )
    ax1.fill_between(x, meanCI, SNICARalbedoLA, alpha=0.2)
    ax1.fill_between(x, SNICARalbedoLA, meanLA, color="k", alpha=0.2)
    ax1.set_xlim(350, 1500), ax1.legend(loc="best")
    ax1.set_ylabel("Albedo"), ax1.set_xlabel("Wavelength (nm)")

    ax2.plot(x, meanCI, linestyle="--", alpha=0.4, label="Clean ice (mean)")
    ax2.plot(x, meanHA, linestyle="-.", alpha=0.4, label="Algal ice (mean)")
    ax2.plot(
        x, SNICARalbedoHA, linestyle="dotted", alpha=0.4, label="Clean ice (model)"
    )
    ax2.fill_between(x, meanCI, SNICARalbedoHA, alpha=0.2)
    ax2.fill_between(x, SNICARalbedoHA, meanHA, color="k", alpha=0.2)
    ax2.set_xlim(350, 1500), ax2.legend(loc="best")
    ax2.set_ylabel("Albedo"), ax2.set_xlabel("Wavelength (nm)")
    fig.tight_layout()
    plt.savefig(str(SAVEPATH + "/BiovsPhysEffect.jpg"), dpi=300)

    # define incoming to calculate broadband albedo
    incoming = xr.open_dataset(
        "/home/joe/Code/BioSNICAR_GO_PY/Data/Mie_files/480band/fsds/swnb_480bnd_sas_clr_SZA60.nc"
    )
    incoming = incoming["flx_frc_sfc"].values
    incoming = incoming[15:230]

    # calculate broadband albedo of each case
    LA_BBA = np.sum(meanLA * incoming) / np.sum(incoming)
    HA_BBA = np.sum(meanHA * incoming) / np.sum(incoming)
    CI_BBA = np.sum(meanCI * incoming) / np.sum(incoming)
    CI2_BBA_LA = np.sum(SNICARalbedoLA * incoming) / np.sum(incoming)
    CI2_BBA_HA = np.sum(SNICARalbedoHA * incoming) / np.sum(incoming)

    # calculate change due to bio/phys as BBA difference
    delAbioLA = CI2_BBA_LA - LA_BBA
    delAphysLA = CI_BBA - CI2_BBA_LA
    delAbioHA = CI2_BBA_HA - HA_BBA
    delAphysHA = CI_BBA - CI2_BBA_HA

    return delAbioLA, delAphysLA, delAbioHA, delAphysHA

def run_best_params(
    SAVEPATH,
    ALL_FIELD_SAMPLES,
    FIELD_DATA_FNAME,
    CI_SITES,
    LA_SITES,
    HA_SITES,
    WEIGHT,
    CLEAN,
):
    """
    function calls out to find_best_params
    """
    ResultArray = np.zeros(shape=(len(ALL_FIELD_SAMPLES), 6))

    for i in np.arange(0, len(ALL_FIELD_SAMPLES), 1):

        fn = ALL_FIELD_SAMPLES[i]
        Results = find_best_params(
            FIELD_DATA_FNAME, fn, CI_SITES, LA_SITES, HA_SITES, WEIGHT, CLEAN
        )
        Results = np.array(Results)
        best_idx = Results[:, 1].argmin()
        best_params = Results[best_idx, 0]
        best_error = Results[best_idx, 1]
        best_dens, best_rds, best_dz, best_alg, best_zen = zip(best_params)
        ResultArray[i, :] = (
            np.array(best_dens),
            np.array(best_rds),
            np.array(best_dz),
            np.array(best_alg),
            np.array(best_zen),
            np.array(best_error),
        )

    Out = pd.DataFrame(
        data=ResultArray, columns=["dens", "rds", "dz", "alg", "zen", "spec_err"]
    )
    Out.index = ALL_FIELD_SAMPLES
    Out.to_csv(str(SAVEPATH + "retrieved_params.csv"))

    return True

def find_best_params(
    FIELD_DATA_FNAME, sampleID, CIsites, LAsites, HAsites, weight, clean=True
):

    """
    This function will return the SNICAR parameter set that provides the
    best approximation to a given field spectrum.

    """

    spectra = pd.read_csv(FIELD_DATA_FNAME)
    spectra = spectra[::10]

    CIspec = spectra[spectra.columns.intersection(CIsites)]
    HAspec = spectra[spectra.columns.intersection(HAsites)]
    LAspec = spectra[spectra.columns.intersection(LAsites)]

    if sampleID == "CImean":
        field_spectrum = CIspec.mean(axis=1)
    elif sampleID == "HAmean":
        field_spectrum = HAspec.mean(axis=1)
    elif sampleID == "LAmean":
        field_spectrum = LAspec.mean(axis=1)
    elif sampleID == "RAIN":
        field_spectrum = spectra["RAIN2"]
    else:
        field_spectrum = spectra[sampleID]

    # calculate 2BDA index of field spectrum
    BDA2idx = np.array(field_spectrum)[36] / np.array(field_spectrum)[31]

    dens = [550, 600, 650, 700, 750, 800, 850]
    dz = [0.02, 0.03, 0.04, 0.05, 0.06, 0.08, 0.1, 0.2]
    rds = [600, 700, 800, 900, 1000]
    alg = [0, 2500, 5000, 7500, 10000, 15000, 20000, 25000, 30000, 35000, 40000, 45000]
    c_factors = [10, 20, 30]
    solzen = [40, 45, 50]

    @dask.delayed
    def run_sims(i, j, k, p, q, z):

        params = c.namedtuple(
            "params",
            "rho_layers, grain_rds, layer_type, dz, mss_cnc_glacier_algae, c_factor_GA, solzen",
        )
        params.rho_layers = [i, i]
        params.grain_rds = [j, j]
        params.layer_type = [1, 1]
        params.dz = [0.001, k]
        params.mss_cnc_glacier_algae = [p, 0]
        params.c_factor_GA = c_factors[q]
        params.solzen = z

        albedo, BBA = call_snicar(params)

        error_vis = np.mean(abs(albedo[15:55] - field_spectrum[0:40]))
        error_nir = np.mean(abs(albedo[55:100] - field_spectrum[40:85]))
        error = ((error_vis * weight) + error_nir) / (1 + weight)

        params = (i, j, k, p, z)
        out = (params, error)

        return out

    Out = []
    # now use the reduced LUT to call snicar and obtain best matching spectrum
    for i in dens:
        for j in rds:
            for k in dz:
                for p in alg:
                    for z in solzen:
                        out = run_sims(i, j, k, p, z)
                        Out.append(out)

    Result = dask.compute(*Out, num_workers=12)

    return Result

def BDA2_of_field_samples():

    """
    2DBA index calculated from field samples after field spectra are averaged over S2 band 4 and 5 wavelengths
    weighted by the sensor spectral response function for each band. The index is then calculated as B5/B4
    and the cell concentration predicted using Wang et al's (2018) conversion equation.

    """

    spectra = pd.read_csv(
        "/home/joe/Code/Remote_Ice_Surface_Analyser/Training_Data/HCRF_master_16171819.csv"
    )

    # reformat LUT: flatten LUT from 3D to 2D array with one column per combination
    # of RT params, one row per wavelength

    responsefunc = pd.read_csv(
        "/home/joe/Code/Remote_Ice_Surface_Analyser/S2SpectralResponse.csv"
    )
    func04 = responsefunc["B4"].loc[
        (responsefunc["SR_WL"] > 650) & (responsefunc["SR_WL"] <= 680)
    ]
    func05 = responsefunc["B5"].loc[
        (responsefunc["SR_WL"] > 698) & (responsefunc["SR_WL"] <= 713)
    ]

    filenames = []
    Idx2DBAList = []
    prd2DBAList = []
    Idx2DBA_S2List = []
    prd2DBA_S2List = []
    Idx2DBA_Ideal_List = []
    prd2DBA_Ideal_List = []

    for i in np.arange(0, len(spectra.columns), 1):

        if i != "Wavelength":

            colname = spectra.columns[i]
            spectrum = np.array(spectra[colname])

            B04 = np.mean(spectrum[300:330] * func04)
            B05 = np.mean(spectrum[348:363] * func05)
            Idx2DBA = spectrum[355] / spectrum[315]
            prd2DBA = 10e-35 * Idx2DBA * np.exp(87.015 * Idx2DBA)
            Idx2DBA_S2 = B05 / B04
            prd2DBA_S2 = 10e-35 * Idx2DBA_S2 * np.exp(87.015 * Idx2DBA_S2)
            Idx2DBA_Ideal = spectrum[360] / spectrum[330]
            prd2DBA_Ideal = 10e-35 * Idx2DBA_Ideal * np.exp(87.015 * Idx2DBA_Ideal)

            filenames.append(colname)
            Idx2DBAList.append(Idx2DBA)
            prd2DBAList.append(prd2DBA)
            Idx2DBA_S2List.append(Idx2DBA_S2)
            prd2DBA_S2List.append(prd2DBA_S2)
            Idx2DBA_Ideal_List.append(Idx2DBA_Ideal)
            prd2DBA_Ideal_List.append(prd2DBA_Ideal)

    Out = pd.DataFrame()
    Out["filename"] = filenames
    Out["2DBAIdx"] = Idx2DBAList
    Out["2DBAPrediction"] = prd2DBAList
    Out["2DBA_S2Idx"] = Idx2DBA_S2List
    Out["2DBA_S2Prediction"] = prd2DBA_S2List
    Out["2DBAIdx_Ideal"] = Idx2DBA_Ideal_List
    Out["2DBAPrediction_Ideal"] = prd2DBA_Ideal_List

    return Out

def compare_predicted_and_measured(SAVEPATH, path_to_metadata):

    ## imports and data organisation
    import numpy as np
    import pandas as pd
    import statsmodels.api as sm

    DF = pd.read_csv(path_to_metadata)

    measured_cells = DF["measured_cells"]
    modelled_cells = DF["algae_cells_inv_model_S2"]
    BDA2_cells = DF["cells_BDA2_centre_wang"]
    BDA2_idx = DF["BDA2_centre"]

    ## regression models

    # Ordinary least squares regression
    model1 = sm.OLS(modelled_cells, measured_cells).fit()
    summary1 = model1.summary()
    test_x = [
        0,
        1000,
        5000,
        7500,
        10000,
        12500,
        15000,
        17500,
        20000,
        25000,
        30000,
        35000,
        40000,
        50000,
    ]
    ypred1 = model1.predict(test_x)

    # regress measured cells against band index
    # use this to give predictive linear model

    BDA2_PredModel = sm.OLS(measured_cells, sm.add_constant(BDA2_idx)).fit()
    BDA2_PredModel_r2 = np.round(BDA2_PredModel.rsquared, 3)
    BDA2_PredModel_y = BDA2_PredModel.predict(sm.add_constant(BDA2_idx))

    # regress BDA2 predicted cells against measured cells
    model2 = sm.OLS(BDA2_PredModel_y, measured_cells).fit()
    summary2 = model2.summary()
    test_x = [
        0,
        1000,
        5000,
        7500,
        10000,
        12500,
        15000,
        17500,
        20000,
        25000,
        30000,
        35000,
        40000,
        50000,
    ]
    ypred2 = model2.predict(test_x)

    # multipanel figure
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 8))

    ax1.plot(
        measured_cells, color="k", marker="x", linestyle="None", label="field-measured"
    )
    ax1.plot(
        modelled_cells,
        color="b",
        marker="o",
        markerfacecolor="None",
        alpha=0.6,
        linestyle="None",
        label="RTM model prediction",
    )
    ax1.plot(
        BDA2_PredModel_y,
        color="r",
        marker="^",
        markerfacecolor="r",
        alpha=0.3,
        linestyle="None",
        label="new 2BDA model prediction",
    )
    ax1.set_ylabel("Algal concentration (cells/mL)")
    ax1.set_xticks(range(len(measured_cells)))
    ax1.set_xticklabels([])
    ax1.legend(loc="upper left")
    ax1.set_xlabel("Individual samples")
    ax1.set_ylim(0, 65000)

    ax2.scatter(
        measured_cells,
        modelled_cells,
        marker="o",
        facecolor="None",
        color="b",
        alpha=0.6,
        label="RTM\nr$^2$ = {}\np = {}".format(
            np.round(model1.rsquared, 3), np.round(model1.pvalues[0], 8)
        ),
    )
    ax2.plot(test_x, ypred1, linestyle="dotted", color="b", alpha=0.6)
    ax2.scatter(
        measured_cells,
        BDA2_PredModel_y,
        marker="^",
        facecolor="r",
        color="r",
        alpha=0.3,
        label="2BDA\nr$^2$ = {}\np = {}".format(
            np.round(model2.rsquared, 3), np.round(model2.pvalues[0], 10)
        ),
    )
    ax2.plot(test_x, ypred2, linestyle="dashed", color="r", alpha=0.3)
    ax2.set_ylabel("Algal concentration,\n cells/mL (field)")
    ax2.set_xlabel("Algal concentration,\n clls/mL (predicted by model)")
    ax2.set_xlim(0, 50000), ax2.set_ylim(0, 60000)

    ax2.legend(loc="upper left")

    fig.tight_layout()

    SAVEPATH = "/home/joe/Code/Remote_Ice_Surface_Analyser/Manuscript/Figures"
    fig.savefig(str(SAVEPATH + "/measured_modelled_algae.png"), dpi=300)

    return


# -*- coding: utf-8 -*-
"""
============================================================================
Generating simple pulses and pulse trains
============================================================================
This example shows how to build and visualize basic types of stimuli such as
:py:class:`~pulse2percept.stimuli.MonophasicPulse`,
:py:class:`~pulse2percept.stimuli.BiphasicPulse` or a
:py:class:`~pulse2percept.stimuli.PulseTrain` for a given implant.

A monophasic pulse has a single phase and can be either anodic (by definition:
has a positive current amplitude) or cathodic (negative current amplitude).

A biphasic pulse is generally charge-balanced for safety reasons (i.e., the
net current must sum to zero over time) and defined as either anodic-first
or cathodic-first.

Multiple pulses can form a pulse train.

"""
# sphinx_gallery_thumbnail_number = 7

##############################################################################
# Simplest stimulus
# ---------------------
# :py:class:`~pulse2percept.stimuli.Stimulus` is the base class to generate
# different types of stimuli. The simplest way to instantiate a Stimulus is
# to pass a scalar value which is interpreted as the current amplitude
# for a single electrode.

# Let's start by importing necessary modules
from pulse2percept.stimuli import (MonophasicPulse, BiphasicPulse,
                                   Stimulus, PulseTrain)

import numpy as np

stim = Stimulus(10)

##############################################################################
# Parameters we don't specify will take on default values. We can inspect
# all current model parameters as follows:

print(stim)

##############################################################################
# This command also reveals a number of other parameters to set, such as:
#
# * ``electrodes``: We can either specify the electrodes in the source
#   or within the stimulus. If none are specified it looks up the source
#   electrode.
#
# * ``metadata``: Optionally we can include metadata to the stimulus we
#   generate as a dictionary.
#
# To change parameter values, either pass them directly to the constructor
# above or set them by hand, like this:

stim.metadata = {'name': 'A simple stimulus', 'date': '2020-01-01'}
stim

##############################################################################
# A monophasic pulse
# --------------------
# We can specify the arguments of the monophasic pulse as follows:

pulse_type = 'anodic'  # anodic: positive amplitude, cathodic: negative
pulse_dur = 4.6 / 1000  # pulse duration in seconds
delay_dur = 10.0 / 1000  # pulse delivered after delay in seconds
stim_dur = 0.5  # stimulus duration in seconds (pulse padded with zeros)
time_step = 0.1 / 1000  # temporal sampling step in seconds

##############################################################################
# The sampling step ``time_step`` defines at which temporal resolution the
# stimulus is resolved. In the above example, the time step is 0.1 ms.
#
# By calling Stimulus with a ``MonophasicPulse`` source, we can generate a
# single pulse:
monophasic_stim = Stimulus(MonophasicPulse(ptype=pulse_type, pdur=pulse_dur,
                                           delay_dur=delay_dur,
                                           stim_dur=stim_dur,
                                           tsample=time_step))
print(monophasic_stim)

##############################################################################
# Here, ``data`` is a 2D NumPy array where rows are electrodes and columns are
# the points in time. Since we did not specify any electrode names in
# ``MonophasicPulse``, the number of electrodes is inferred from the input
# source type. There is only one row in the above example, denoting a single
# electrode.
#
# By default, the :py:class:`~pulse2percept.stimuli.MonophasicPulse` object
# automatically assumes a current amplitude of 1 uA.

##############################################################################
# We can visualize the generated pulse using Matplotlib:

import matplotlib.pyplot as plt
fig, ax = plt.subplots(figsize=(8, 5))
ax.plot(monophasic_stim.time, monophasic_stim.data[0, :])
ax.set_xlabel('Time (s)')
ax.set_ylabel('Amplitude ($\mu$A)')

###############################################################################
# A biphasic pulse
# ------------------
# Similarly, we can generate a biphasic pulse by changing the source of the
# stimulus to :py:class:`~pulse2percept.stimuli.BiphasicPulse`. This time
# parameter ``ptype`` can either be 'anodicfirst' or 'cathodicfirst'.

# set relevant parameters
pulse_type = 'cathodicfirst'
biphasic_stim = Stimulus(BiphasicPulse(ptype=pulse_type, pdur=pulse_dur,
                                       tsample=time_step))

###############################################################################
# If we visualize this stimulus, we can see the difference between a monophasic
# and biphasic pulse:

# Create a figure with two subplots
fig, axes = plt.subplots(1, 2, figsize=(8, 5))

# First, plot monophasic pulse
axes[0].plot(monophasic_stim.time, monophasic_stim.data[0])
ax.set_xlabel('Time (s)')
ax.set_ylabel('Amplitude ($\mu$A)')

# Second, plot biphasic pulse
axes[1].plot(biphasic_stim.time, biphasic_stim.data[0])
ax.set_xlabel('Time (s)')
ax.set_ylabel('Amplitude ($\mu$A)')

###############################################################################
# Changing pulse amplitude
# ----------------------------------
# For any given pulse, we can modify the amplitude by indexing into the ``data``
# row that corresponds to the desired electrode. In the above example, we only
# have one electrode (index 0).
# Let's say we want the amplitude of the monophasic pulse to be 10 micro amps.
# We have two options: either change the values of the ``data`` array directly:

# get the data structure by indexing the electrode at 0
monophasic_stim.data[0] = 10 * monophasic_stim.data[0]
print(monophasic_stim)

###############################################################################
# Or we can create a NumPy array and assign that to the data structure of the
# stimulus:

# recreate the same stimulus with an amplitude 1 microAmps.
monophasic_stim = Stimulus(MonophasicPulse(ptype='anodic', pdur=pulse_dur,
                                           delay_dur=delay_dur,
                                           stim_dur=stim_dur,
                                           tsample=time_step))
monophasic_stim.data[0] = 10 * np.ones_like(monophasic_stim.data[0])
print(monophasic_stim)

###############################################################################
# Similarly, let's say we want the cathodic part of the biphasic pulse to be -5
# micro amps, and the anodic part to be +20 micro amps (note that this stimulus
# wouldn't be charge-balanced).
#
# We first need to find the halfway point where the current switches from
# cathodic to anodic. To do that we first get the length of the pulse by
# indexing the single electrode at 0
length = len(biphasic_stim.data[0])
print(length)

# Find the halfway where cathodic turns into anodic pulse
half = int(len(biphasic_stim.data[0]) / 2)
print("Halfway index is", half)

# change the first half of the pulse to be 5 times larger
biphasic_stim.data[0][0:half] = 5 * biphasic_stim.data[0][0:half]

# change the second half to be 20 times larger
biphasic_stim.data[0][half:length] = 20 * biphasic_stim.data[0][half:length]

###############################################################################
# Let's plot the monophasic and biphasic pulses again:

# Create a figure with two subplots
fig, axes = plt.subplots(ncols=2, figsize=(8, 5))

# First, plot monophasic pulse
axes[0].plot(monophasic_stim.time, monophasic_stim.data[0])
axes[0].set_xlabel('Time (s)')
axes[0].set_ylabel('Amplitude ($\mu$A)')
# Second, plot biphasic pulse
axes[1].plot(biphasic_stim.time, biphasic_stim.data[0])
axes[1].set_xlabel('Time (s)')
axes[1].set_ylabel('Amplitude ($\mu$A)')
fig.tight_layout()

###############################################################################
# Generating standard pulse trains
# ----------------------------------
# The easiest way to generate a pulse train is to use the
# :py:class:`~pulse2percept.stimuli.PulseTrain` object, which allows for
# various stimulus attributes to be specified:

time_step = 0.1 / 1000  # temporal sampling in seconds
freq = 20  # frequency in Hz
amp = 100  # maximum amplitude of the pulse train in microAmps
dur = 0.2  # total duration of the pulse train in seconds
pulse_type = 'cathodicfirst'  # whether the first phase is positive or negative
pulse_order = 'gapfirst'  # whether the train starts with gap or a pulse.

# Define the pulse train with given parameters
ptrain = PulseTrain(tsample=time_step,
                    freq=freq,
                    dur=dur,
                    amp=amp,
                    pulsetype=pulse_type,
                    pulseorder=pulse_order)

# Create a new stimulus where the pulse train is the source
ptrain_stim = Stimulus(ptrain)

# Visualize:
fig, ax = plt.subplots(figsize=(8, 5))
ax.plot(ptrain_stim.time, ptrain_stim.data[0, :])
ax.set_xlabel('Time (s)')
ax.set_ylabel('Amplitude ($\mu$A)')

###############################################################################
# Alternatively, we are free to specify a discrete set of points in time and
# the current amplitude we would like to apply at those times.
#
# It is important to note that the :py:class:`~pulse2percept.stimuli.Stimulus`
# object will linearly interpolate between specified time points.
# For example, the following generates a simple sawtooth stimulus:

stim = Stimulus([[0, -10, 10, -10, 10, -10, 0]],
                time=[0, 0.1, 0.3, 0.5, 0.7, 0.9, 1.0])
fig, ax = plt.subplots(figsize=(8, 5))
ax.plot(stim.time, stim.data[0, :])
ax.set_xlabel('Time (s)')
ax.set_ylabel('Amplitude ($\mu$A)')

##############################################################################
# For a biphasic pulse, we need to specify both the rising edge (low-to-high
# transition) and falling edge (high-to-low transition) of the signal:

stim = Stimulus([[0, 0, 10, 10,  0, 0]],
                time=[0, 0.1, 0.1, 0.2, 0.2, 1.0])
fig, ax = plt.subplots(figsize=(8, 5))
ax.plot(stim.time, stim.data[0, :])
ax.set_xlabel('Time (s)')
ax.set_ylabel('Amplitude ($\mu$A)')

##############################################################################
# We can thus generate arbitrarily complex stimuli:

stim = Stimulus([[0, 0, 20, 20, -5, -5, 0, 0, 0, 20, 20, -5, -5, 0, 0]],
                time=[0, 0.1, 0.1, 0.2, 0.2, 0.6, 0.6, 1.0, 1.1, 1.1, 1.2, 1.2, 1.6, 1.6, 2.0])
fig, ax = plt.subplots(figsize=(8, 5))
ax.plot(stim.time, stim.data[0, :])
ax.set_xlabel('Time (s)')
ax.set_ylabel('Amplitude ($\mu$A)')


from __future__ import (absolute_import, division, print_function, unicode_literals)
import numpy as np
import multiprocessing
from functools import partial

from ..pair_counters.npairs_xy_z import _npairs_xy_z_process_args
from ..pair_counters.mesh_helpers import _set_approximate_cell_sizes, _cell1_parallelization_indices
from ..pair_counters.rectangular_mesh import RectangularDoubleMesh
from .velocity_marked_npairs_3d import _velocity_marked_npairs_3d_process_weights
from .engines import velocity_marked_npairs_xy_z_engine

__author__ = ('Duncan Campbell', 'Andrew Hearin')

__all__ = ('velocity_marked_npairs_xy_z', )

def velocity_marked_npairs_xy_z(sample1, sample2, rp_bins, pi_bins, period=None,
        weights1=None, weights2=None, weight_func_id=1, num_threads=1,
        approx_cell1_size=None, approx_cell2_size=None):
    r"""
    Calculate the number of velocity weighted pairs
    with separations greater than or equal to
    :math:`r_{\perp}` and :math:`r_{\parallel}`, :math:`W(>r_{\perp},>r_{\parallel})`.

    :math:`r_{\perp}` and :math:`r_{\parallel}` are defined wrt the z-direction.

    The weight given to each pair is determined by the weights for a pair,
    :math:`w_1`, :math:`w_2`, and a user-specified "velocity weighting function", indicated
    by the ``weight_func_id`` parameter, :math:`f(w_1,w_2)`.

    Parameters
    ----------
    sample1 : array_like
        Npts1 x 3 numpy array containing 3-D positions of points.
        See the :ref:`mock_obs_pos_formatting` documentation page, or the
        Examples section below, for instructions on how to transform
        your coordinate position arrays into the
        format accepted by the ``sample1`` and ``sample2`` arguments.
        Length units are comoving and assumed to be in Mpc/h, here and throughout Halotools.

    sample2 : array_like
        Npts2 x 3 array containing 3-D positions of points.

    rp_bins : array_like
        array of boundaries defining the radial bins perpendicular to the LOS in which
        pairs are counted.
        Length units are comoving and assumed to be in Mpc/h, here and throughout Halotools.

    pi_bins : array_like
        array of boundaries defining the bins parallel to the LOS in which
        pairs are counted.
        Length units are comoving and assumed to be in Mpc/h, here and throughout Halotools.

    period : array_like, optional
        Length-3 sequence defining the periodic boundary conditions
        in each dimension. If you instead provide a single scalar, Lbox,
        period is assumed to be the same in all Cartesian directions.
        If set to None (the default option), PBCs are set to infinity.
        Length units are comoving and assumed to be in Mpc/h, here and throughout Halotools.

    weights1 : array_like, optional
        Either a 1-D array of length *Npts1*, or a 2-D array of length *Npts1* x *N_weights*,
        containing the weights used for the weighted pair counts. If this parameter is
        None, the weights are set to np.ones(*(Npts1,N_weights)*).

    weights2 : array_like, optional
        Either a 1-D array of length *Npts2*, or a 2-D array of length *Npts2* x *N_weights*,
        containing the weights used for the weighted pair counts. If this parameter is
        None, the weights are set to np.ones(*(Npts2,N_weights)*).

    weight_func_id : int, optional
        velocity weighting function integer ID. Each weighting function requires a specific
        number of weights per point, *N_weights*.  See the Notes for a description of
        available weighting functions.

    num_threads : int, optional
        Number of threads to use in calculation, where parallelization is performed
        using the python ``multiprocessing`` module. Default is 1 for a purely serial
        calculation, in which case a multiprocessing Pool object will
        never be instantiated. A string 'max' may be used to indicate that
        the pair counters should use all available cores on the machine.

    approx_cell1_size : array_like, optional
        Length-3 array serving as a guess for the optimal manner by how points
        will be apportioned into subvolumes of the simulation box.
        The optimum choice unavoidably depends on the specs of your machine.
        Default choice is to use Lbox/10 in each dimension,
        which will return reasonable result performance for most use-cases.
        Performance can vary sensitively with this parameter, so it is highly
        recommended that you experiment with this parameter when carrying out
        performance-critical calculations.

    approx_cell2_size : array_like, optional
        Analogous to ``approx_cell1_size``, but for sample2.  See comments for
        ``approx_cell1_size`` for details.

    Returns
    -------
    w1N_pairs : numpy.array
        2-D array of shape *(Nrp_bins,Npi_bins)* containing the weighted number counts
        of pairs. The exact values depend on ``weight_func_id``
        (which weighting function was chosen).

    w2N_pairs : numpy.array
        2-D array of shape *(Nrp_bins,Npi_bins)* containing the weighted number counts
        of pairs. The exact values depend on ``weight_func_id``
        (which weighting function was chosen).

    w3N_pairs : numpy.array
        2-D array of shape *(Nrp_bins,Npi_bins)* containing the weighted number counts
        of pairs. The exact values depend on ``weight_func_id``
        (which weighting function was chosen).

    Examples
    --------
    For demonstration purposes we will work with
    halos in the `~halotools.sim_manager.FakeSim`.

    >>> from halotools.sim_manager import FakeSim
    >>> halocat = FakeSim()

    >>> x = halocat.halo_table['halo_x']
    >>> y = halocat.halo_table['halo_y']
    >>> z = halocat.halo_table['halo_z']

    We transform our *x, y, z* points into the array shape used by the pair-counter by
    taking the transpose of the result of `numpy.vstack`. This boilerplate transformation
    is used throughout the `~halotools.mock_observables` sub-package:

    >>> sample1 = np.vstack((x,y,z)).T

    We will do the same to get a random set of velocities.

    >>> vx = halocat.halo_table['halo_vx']
    >>> vy = halocat.halo_table['halo_vy']
    >>> vz = halocat.halo_table['halo_vz']
    >>> velocities = np.vstack((x,y,z,vx,vy,vz)).T

    >>> rp_bins = np.logspace(-2,-1,10)
    >>> pi_bins = np.linspace(0, 10, 5)
    >>> result = velocity_marked_npairs_xy_z(sample1, sample1, rp_bins, pi_bins, period=halocat.Lbox, weights1=velocities, weights2=velocities,)

    """
    result = _npairs_xy_z_process_args(sample1, sample2, rp_bins, pi_bins, period,
            num_threads, approx_cell1_size, approx_cell2_size)
    x1in, y1in, z1in, x2in, y2in, z2in = result[0:6]
    rp_bins, pi_bins, period, num_threads, PBCs, approx_cell1_size, approx_cell2_size = result[6:]
    xperiod, yperiod, zperiod = period

    rp_max = np.max(rp_bins)
    pi_max = np.max(pi_bins)
    search_xlength, search_ylength, search_zlength = rp_max, rp_max, pi_max

    # Process the input weights and with the helper function
    weights1, weights2 = (
        _velocity_marked_npairs_3d_process_weights(sample1, sample2,
            weights1, weights2, weight_func_id))

    # Compute the estimates for the cell sizes
    approx_cell1_size, approx_cell2_size = (
        _set_approximate_cell_sizes(approx_cell1_size, approx_cell2_size, period)
        )
    approx_x1cell_size, approx_y1cell_size, approx_z1cell_size = approx_cell1_size
    approx_x2cell_size, approx_y2cell_size, approx_z2cell_size = approx_cell2_size

    # Build the rectangular mesh
    double_mesh = RectangularDoubleMesh(x1in, y1in, z1in, x2in, y2in, z2in,
        approx_x1cell_size, approx_y1cell_size, approx_z1cell_size,
        approx_x2cell_size, approx_y2cell_size, approx_z2cell_size,
        search_xlength, search_ylength, search_zlength, xperiod, yperiod, zperiod, PBCs)

    # Create a function object that has a single argument, for parallelization purposes
    engine = partial(velocity_marked_npairs_xy_z_engine, double_mesh,
        x1in, y1in, z1in, x2in, y2in, z2in,
        weights1, weights2, weight_func_id, rp_bins, pi_bins)

    # Calculate the cell1 indices that will be looped over by the engine
    num_threads, cell1_tuples = _cell1_parallelization_indices(
        double_mesh.mesh1.ncells, num_threads)

    if num_threads > 1:
        pool = multiprocessing.Pool(num_threads)
        result = np.array(pool.map(engine, cell1_tuples))
        counts1, counts2, counts3 = result[:, 0], result[:, 1], result[:, 2]
        counts1 = np.sum(counts1, axis=0)
        counts2 = np.sum(counts2, axis=0)
        counts3 = np.sum(counts3, axis=0)
        pool.close()
    else:
        counts1, counts2, counts3 = np.array(engine(cell1_tuples[0]))

    return counts1, counts2, counts3


from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple, Union, cast

import numpy as np
import tensorflow as tf
from srl.base.rl.algorithms.continuous_action import ContinuousActionConfig, ContinuousActionWorker
import tensorflow.keras as keras
import tensorflow.keras.layers as kl
from srl.base.define import RLObservationType
from srl.base.env.base import EnvRun
from srl.base.rl.algorithms.continuous_action import ContinuousActionConfig, ContinuousActionWorker
from srl.base.rl.base import RLParameter, RLTrainer
from srl.base.rl.registration import register
from srl.base.rl.remote_memory import ExperienceReplayBuffer
from srl.rl.functions.common_tf import compute_logprob_sgp
from srl.rl.functions.model import ImageLayerType, create_input_layers_one_sequence

"""
DDPG
    Replay buffer       : o
    Target Network(soft): o
    Target Network(hard): o
    Add action noise    : x
TD3
    Clipped Double Q learning : o
    Target Policy Smoothing   : x
    Delayed Policy Update     : x
SAC
    Squashed Gaussian Policy: o
"""

# ------------------------------------------------------
# config
# ------------------------------------------------------
@dataclass
class Config(ContinuousActionConfig):

    # model
    policy_hidden_layer_sizes: Tuple[int, ...] = (64, 64, 64)
    q_hidden_layer_sizes: Tuple[int, ...] = (64, 64, 64)
    activation: str = "relu"
    image_layer_type: ImageLayerType = ImageLayerType.DQN

    gamma: float = 0.9  # 割引率
    lr: float = 0.005  # 学習率
    soft_target_update_tau: float = 0.02
    hard_target_update_interval: int = 100

    batch_size: int = 32
    capacity: int = 10_000
    memory_warmup_size: int = 500

    def __post_init__(self):
        super().__init__()

    @property
    def observation_type(self) -> RLObservationType:
        return RLObservationType.CONTINUOUS

    @staticmethod
    def getName() -> str:
        return "SAC"

    def assert_params(self) -> None:
        super().assert_params()
        assert self.memory_warmup_size < self.capacity
        assert self.batch_size < self.memory_warmup_size

register(
    Config,
    __name__ + ":RemoteMemory",
    __name__ + ":Parameter",
    __name__ + ":Trainer",
    __name__ + ":Worker",
)

# ------------------------------------------------------
# RemoteMemory
# ------------------------------------------------------
class RemoteMemory(ExperienceReplayBuffer):
    def __init__(self, *args):
        super().__init__(*args)
        self.config = cast(Config, self.config)

        self.init(self.config.capacity)

# ------------------------------------------------------
# network
# ------------------------------------------------------
class _PolicyModel(keras.Model):
    def __init__(self, config: Config):
        super().__init__()

        in_state, c = create_input_layers_one_sequence(
            config.observation_shape,
            config.env_observation_type,
            config.image_layer_type,
        )

        # --- hidden layer
        for h in config.policy_hidden_layer_sizes:
            c = kl.Dense(h, activation=config.activation, kernel_initializer="he_normal")(c)
        c = kl.LayerNormalization()(c)  # 勾配爆発抑制用?

        # --- out layer
        pi_mean = kl.Dense(
            config.action_num,
            activation="linear",
            kernel_initializer="truncated_normal",
            bias_initializer="truncated_normal",
        )(c)
        pi_stddev = kl.Dense(
            config.action_num,
            activation="linear",
            kernel_initializer="truncated_normal",
            bias_initializer="truncated_normal",
        )(c)
        self.model = keras.Model(in_state, [pi_mean, pi_stddev])

        # 重みを初期化
        dummy_state = np.zeros(shape=(1,) + config.observation_shape, dtype=np.float32)
        action, mean, stddev, action_org = self(dummy_state)
        assert mean.shape == (1, config.action_num)
        assert stddev.shape == (1, config.action_num)

    @tf.function
    def call(self, state):
        mean, stddev = self.model(state)

        # σ > 0
        stddev = tf.exp(stddev)

        # Reparameterization trick
        normal_random = tf.random.normal(mean.shape, mean=0.0, stddev=1.0)
        action_org = mean + stddev * normal_random

        # Squashed Gaussian Policy
        action = tf.tanh(action_org)

        return action, mean, stddev, action_org

class _DualQNetwork(keras.Model):
    def __init__(self, config: Config):
        super().__init__()

        in_state, c = create_input_layers_one_sequence(
            config.observation_shape,
            config.env_observation_type,
            config.image_layer_type,
        )
        in_action = kl.Input(shape=(config.action_num,))
        c = kl.Concatenate()([c, in_action])

        for h in config.q_hidden_layer_sizes:
            c = kl.Dense(h, activation=config.activation, kernel_initializer="he_normal")(c)
        c = kl.LayerNormalization()(c)  # 勾配爆発抑制用?

        # q1
        q1 = kl.Dense(
            1, activation="linear", kernel_initializer="truncated_normal", bias_initializer="truncated_normal"
        )(c)

        # q2
        q2 = kl.Dense(
            1, activation="linear", kernel_initializer="truncated_normal", bias_initializer="truncated_normal"
        )(c)

        # out layer
        self.model = keras.Model([in_state, in_action], [q1, q2])

        # 重みを初期化
        dummy_state = np.zeros(shape=(1,) + config.observation_shape, dtype=np.float32)
        dummy_action = np.zeros(shape=(1, config.action_num), dtype=np.float32)
        _q1, _q2 = self(dummy_state, dummy_action)
        assert _q1.shape == (1, 1)
        assert _q2.shape == (1, 1)

    def call(self, state, action):
        return self.model([state, action])

# ------------------------------------------------------
# Parameter
# ------------------------------------------------------
class Parameter(RLParameter):
    def __init__(self, *args):
        super().__init__(*args)
        self.config = cast(Config, self.config)

        self.policy = _PolicyModel(self.config)
        self.q_online = _DualQNetwork(self.config)
        self.q_target = _DualQNetwork(self.config)

    def restore(self, data: Any) -> None:
        self.policy.set_weights(data[0])
        self.q_online.set_weights(data[1])
        self.q_target.set_weights(data[1])

    def backup(self) -> Any:
        return [
            self.policy.get_weights(),
            self.q_online.get_weights(),
        ]

    def summary(self):
        self.policy.model.summary()
        self.q_online.model.summary()

# ------------------------------------------------------
# Trainer
# ------------------------------------------------------
class Trainer(RLTrainer):
    def __init__(self, *args):
        super().__init__(*args)
        self.config = cast(Config, self.config)
        self.parameter = cast(Parameter, self.parameter)
        self.remote_memory = cast(RemoteMemory, self.remote_memory)

        self.train_count = 0

        self.q_optimizer = keras.optimizers.Adam(learning_rate=self.config.lr)
        self.policy_optimizer = keras.optimizers.Adam(learning_rate=self.config.lr)
        self.alpha_optimizer = keras.optimizers.Adam(learning_rate=self.config.lr)

        # エントロピーαの目標値、-1×アクション数が良いらしい
        self.target_entropy = -1 * self.config.action_num

        # エントロピーα自動調整用
        self.log_alpha = tf.Variable(0.5, dtype=tf.float32)

    def get_train_count(self):
        return self.train_count

    def train(self):
        if self.remote_memory.length() < self.config.memory_warmup_size:
            return {}
        batchs = self.remote_memory.sample(self.config.batch_size)

        states = []
        actions = []
        n_states = []
        rewards = []
        dones = []
        for b in batchs:
            states.append(b["state"])
            actions.append(b["action"])
            n_states.append(b["next_state"])
            rewards.append(b["reward"])
            dones.append(b["done"])
        states = np.asarray(states)
        n_states = np.asarray(n_states)
        actions = np.asarray(actions)
        dones = np.asarray(dones).reshape((-1, 1))
        rewards = np.asarray(rewards).reshape((-1, 1))

        # 方策エントロピーの反映率αを計算
        alpha = tf.math.exp(self.log_alpha)

        # ポリシーより次の状態のアクションを取得
        n_actions, n_means, n_stddevs, n_action_orgs = self.parameter.policy(n_states)
        # 次の状態のアクションのlogpiを取得(Squashed Gaussian Policy時)
        n_logpi = compute_logprob_sgp(n_means, n_stddevs, n_action_orgs)

        # 2つのQ値から小さいほうを採用(Clipped Double Q learning)して、
        # Q値を計算 : reward if done else (reward + gamma * n_qval) - (alpha * H)
        n_q1, n_q2 = self.parameter.q_target(n_states, n_actions)
        q_vals = rewards + (1 - dones) * self.config.gamma * tf.minimum(n_q1, n_q2) - (alpha * n_logpi)

        # --- Qモデルの学習
        with tf.GradientTape() as tape:
            q1, q2 = self.parameter.q_online(states, actions)
            loss1 = tf.reduce_mean(tf.square(q_vals - q1))
            loss2 = tf.reduce_mean(tf.square(q_vals - q2))
            q_loss = (loss1 + loss2) / 2

        grads = tape.gradient(q_loss, self.parameter.q_online.trainable_variables)
        self.q_optimizer.apply_gradients(zip(grads, self.parameter.q_online.trainable_variables))

        # --- ポリシーの学習
        with tf.GradientTape() as tape:
            # アクションを出力
            selected_actions, means, stddevs, action_orgs = self.parameter.policy(states)

            # logπ(a|s) (Squashed Gaussian Policy)
            logpi = compute_logprob_sgp(means, stddevs, action_orgs)

            # Q値を出力、小さいほうを使う
            q1, q2 = self.parameter.q_online(states, selected_actions)
            q_min = tf.minimum(q1, q2)

            # alphaは定数扱いなので勾配が流れないようにする
            policy_loss = q_min - (tf.stop_gradient(alpha) * logpi)

            policy_loss = -tf.reduce_mean(policy_loss)  # 最大化

        grads = tape.gradient(policy_loss, self.parameter.policy.trainable_variables)
        self.policy_optimizer.apply_gradients(zip(grads, self.parameter.policy.trainable_variables))

        # --- 方策エントロピーαの自動調整
        _, means, stddevs, action_orgs = self.parameter.policy(states)
        logpi = compute_logprob_sgp(means, stddevs, action_orgs)

        with tf.GradientTape() as tape:
            entropy_diff = -logpi - self.target_entropy
            log_alpha_loss = tf.reduce_mean(tf.exp(self.log_alpha) * entropy_diff)

        grad = tape.gradient(log_alpha_loss, self.log_alpha)
        self.alpha_optimizer.apply_gradients([(grad, self.log_alpha)])

        # --- soft target update
        self.parameter.q_target.set_weights(
            (1 - self.config.soft_target_update_tau) * np.array(self.parameter.q_target.get_weights(), dtype=object)
            + (self.config.soft_target_update_tau) * np.array(self.parameter.q_online.get_weights(), dtype=object)
        )

        # --- hard target sync
        if self.train_count % self.config.hard_target_update_interval == 0:
            self.parameter.q_target.set_weights(self.parameter.q_online.get_weights())

        self.train_count += 1
        return {
            "q_loss": q_loss.numpy(),
            "policy_loss": policy_loss.numpy(),
            "alpha_loss": log_alpha_loss.numpy(),
        }

# ------------------------------------------------------
# Worker
# ------------------------------------------------------
class Worker(ContinuousActionWorker):
    def __init__(self, *args):
        super().__init__(*args)
        self.config = cast(Config, self.config)
        self.parameter = cast(Parameter, self.parameter)
        self.remote_memory = cast(RemoteMemory, self.remote_memory)

    def call_on_reset(self, state: np.ndarray) -> None:
        self.state = state

    def call_policy(self, state: np.ndarray) -> List[float]:
        self.state = state
        action, mean, _, _ = self.parameter.policy(state.reshape(1, -1))

        if self.training:
            action = action.numpy()[0]
        else:
            # テスト時は平均を使う
            action = mean.numpy()[0]

        # Squashed Gaussian Policy (-1, 1) -> (action range)
        action = (action + 1) / 2
        action = self.config.action_low + action * (self.config.action_high - self.config.action_low)

        self.action = action
        return self.action

    def call_on_step(
        self,
        next_state: np.ndarray,
        reward: float,
        done: bool,
    ) -> Dict[str, Union[float, int]]:
        if not self.training:
            return {}

        batch = {
            "state": self.state,
            "action": self.action,
            "next_state": next_state,
            "reward": reward,
            "done": done,
        }
        self.remote_memory.add(batch)

        return {}

    def call_render(self, env: EnvRun) -> None:
        state = self.state.reshape(1, -1)
        action = np.asarray([self.action])
        _, mean, stddev, _ = self.parameter.policy(state)
        mean = mean.numpy()[0][0]
        stddev = stddev.numpy()[0][0]
        q1, q2 = self.parameter.q_online(state, action)
        q1 = q1.numpy()[0][0]
        q2 = q2.numpy()[0][0]

        print(f"mean {mean:.5f}, stddev {stddev:.5f}")
        print(f"q1   {q1:.5f}, q2    {q2:.5f}")


import sys, re, glob
import numpy as np

from cctk import GaussianFile, Molecule
from cctk import parse_gaussian as parse

filenames = sys.argv[1]
info = []

for filename in sorted(glob.glob(filenames, recursive=True)):
    try:
        output_file, lines = GaussianFile.read_file(filename, return_lines=True)
    except:
        continue

    success = "NO"
    if output_file.success:
        success = output_file.success
    else:
        continue

    energy = output_file.energies[-1]
    iters = len(output_file.energies)
    mol = output_file.get_molecule()

    cation_anion_dist = 0
    mul_q = 0
    hir_q = 0
    if 29 in mol.atomic_numbers:
        mul_q = float(parse.find_parameter(lines, "    52  Cu", 4, 2)[-1])
        hir_q = float(parse.find_parameter(lines, "    52  Cu", 8, 2)[-1])

        if 16 in mol.atomic_numbers:
            cation_anion_dist = mol.get_distance(52, 54)
        elif 51 in mol.atomic_numbers:
            cation_anion_dist = mol.get_distance(52, 53)

    imaginaries = "--"
    try:
        if output_file.num_imaginaries() > 0:
            if output_file.num_imaginaries() > 1:
                imaginaries = ", ".join(output_file.imaginaries())
            else:
                imaginaries = output_file.imaginaries()[0]
    except:
        #### Will raise ValueError if job is not of type "FREQ"
        pass

    info.append([filename, energy, energy * 627.509, iters, mul_q, hir_q, success, imaginaries, cation_anion_dist])

if len(info) > 0:
    min_energy = np.min([x[2] for x in info])
    def adjust_energy(row):
        if row[2] < 0:
            row[2] = row[2] - min_energy
        return row

    info = list(map(adjust_energy, info))

    print("{0},{1},{2},{3},{4},{5},{6},{7},{8}".format(
        "File", "Energy (Hartree)", "Rel Energy (kcal)", "Iterations", "Mulliken Q", "Hirshfeld Q", "Success?", "Imaginaries?","X-Cu Distance"
    ))

    for row in info:
        print("{0},{1:.4f},{2:.2f},{3},{4:.4f},{5:.4f},{6},{7:},{8:.2f}".format(*row))