#!/usr/bin/env python

# Copyright 2017-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file 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.

import os
import json
import sys
from io import BytesIO
import argparse
import numpy as np
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
from scipy.cluster import hierarchy
from scipy.spatial.distance import squareform, pdist
import task_similarity
import glob
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns

CATEGORIES_JSON_FILE = 'inat2018/categories.json'
ICONS_PATH = './static/iconic_taxa/'

parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
parser.add_argument('root', default='multirun/variational', type=str)
parser.add_argument('--data-root', default='~/data', type=str)
parser.add_argument('--distance', default='cosine', type=str,
                    help='distance to use')
parser.add_argument('--max-tasks', default=120, type=int,
                    help='number of tasks to consider')
args = parser.parse_args()

# Assumes task IDs are mapped as follows:
# CUB: 0-24 (0-9 are orders, 10-24 are Passeriformes families)
# iNat2018: 25-231
# As in cub_inat2018 in datasets.py
CUB = 'CUB'
INAT = 'iNat'
CUB_NUM_TASKS = 25

ADDITIONAL_TAXONOMY_DATA = [
    {
        'kingdom': 'Animalia ',
        'supercategory': 'Animalia ',
        'phylum': 'Chordata',
        'class': 'Aves',
        'order': 'Apodiformes',
    }
]

def invert_icon(img):
    img = (1. - img)
    return img

def get_image(e):
    base = os.path.join(ICONS_PATH, "{}-200px.png")
    possible_names = [base.format(x) for x in
                      [e.meta.get('class'), e.meta.get('phylum'), e.meta.get('kingdom'), 'unknown']]
    for filename in possible_names:
        if os.path.exists(filename):
            img = plt.imread(filename, format='png')
            return invert_icon(img) if e.dataset == CUB else img
    raise FileNotFoundError()

def average_top_k_tax_distance(distance_matrix, from_embeddings, to_embeddings=None, k=2):
    assert k > 0, k

    if to_embeddings is None:
        to_embeddings = from_embeddings

    assert distance_matrix.shape[0] == len(from_embeddings)
    assert distance_matrix.shape[1] == len(to_embeddings)

    tax_distance = []
    for i in range(len(from_embeddings)):
        top_matches = distance_matrix[i].argsort()[:k]
        tax_distance.append(
            np.mean([taxonomy_distance(from_embeddings[i], to_embeddings[j]) for j in top_matches])
        )
    return np.mean(tax_distance)

def plot_changing_k(ax, distance_matrix, from_embeddings, to_embeddings, **kwargs):
    x = [1, 3, 5, 10, 20, 30, 50, 100, 200, len(to_embeddings)]
    x = [v for v in x if v <= len(to_embeddings)]
    y = []
    for k in x:
        y.append(average_top_k_tax_distance(distance_matrix, from_embeddings, to_embeddings, k=k))
    ax.plot(x, y, **kwargs)
    ax.set_xlabel('Size k of neighborhood')
    ax.set_ylabel('Avg. top-k tax. distance')

def sort_distance_matrix(distance_matrix, embeddings, names, method='complete'):
    assert method in ['ward', 'single', 'average', 'complete']
    np.fill_diagonal(distance_matrix, 0.)
    cond_distance_matrix = squareform(distance_matrix, checks=False)
    linkage_matrix = hierarchy.linkage(cond_distance_matrix, method='complete', optimal_ordering=True)
    res_order = hierarchy.leaves_list(linkage_matrix)
    distance_matrix = distance_matrix[res_order][:, res_order]
    embeddings = [embeddings[i] for i in res_order]
    names = [names[i] for i in res_order]
    np.fill_diagonal(distance_matrix, np.nan)
    return distance_matrix, embeddings, names, res_order

def draw_figure_to_plt(distance_matrix, embeddings, names, label_size=14):
    fig = plt.figure(figsize=(15 / 25. * len(embeddings), 15 / 25. * len(embeddings)))
    ax = plt.gca()

    plt.imshow(distance_matrix, cmap='viridis_r')
    ax.set_xticks(np.arange(len(embeddings)))
    ax.set_yticks(np.arange(len(embeddings)))

    ax.set_xticklabels(names)
    ax.set_yticklabels(names)
    plt.setp(ax.get_xticklabels(), rotation=45, ha="right",
             rotation_mode="anchor")

    try:
        for i, e in enumerate(embeddings):
            arr_img = get_image(e)
            imagebox = OffsetImage(arr_img, zoom=0.18)
            imagebox.image.axes = ax
            xy = (i, i)
            ab = AnnotationBbox(imagebox, xy, frameon=False)
            ax.add_artist(ab)
    except FileNotFoundError:
        print("Could not find an icon for a taxonomy entry. Have you downloaded the iconic_taxa directory in ./static?")

    plt.tick_params(axis='both', which='major', labelsize=label_size)
    plt.tight_layout()

def taxonomy_distance(e0, e1):
    for i, k in enumerate(['order', 'class', 'phylum', 'kingdom']):
        if e0.meta[k] == e1.meta[k]:
            return i
    return i + 1

def add_class_information(embeddings):
    # load taxonomy
    with open(os.path.join(args.data_root, CATEGORIES_JSON_FILE), 'r') as f:
        categories = json.load(f)
    categories.extend(ADDITIONAL_TAXONOMY_DATA)

    category_map = {c['order']: c for c in categories}
    category_map.update({c['family']: c for c in categories if 'family' in c})

    for e in embeddings:
        c = category_map[e.task_name]
        e.meta['order'] = c['order'].lower()
        e.meta['class'] = c['class'].lower()
        e.meta['phylum'] = c['phylum'].lower()
        e.meta['kingdom'] = c['kingdom'].lower()
        e.meta['supercategory'] = c['supercategory'].lower()

def main():
    os.makedirs('./plots', exist_ok=True)

    files = glob.glob(os.path.join(args.root, '*', 'embedding.p'))

    # get embeddings
    embeddings = [task_similarity.load_embedding(file) for file in files]
    embeddings.sort(key=lambda x: x.meta['dataset']['task_id'])
    for e in embeddings:
        e.task_id = e.meta['dataset']['task_id']
        e.task_name = e.meta['task_name']
        e.dataset = CUB if e.task_id < CUB_NUM_TASKS else INAT

    # get distance matrix
    distance_matrix = task_similarity.pdist(embeddings, distance=args.distance)
    add_class_information(embeddings)

    # construct names to display on plot
    for e in embeddings:
        assert hasattr(e, 'task_name')
        assert 'class' in e.meta, e.task_name
        assert hasattr(e, 'dataset'), e.task_name

    names = [f"[{e.dataset}] {e.task_name} ({e.meta['class']})" if 'order' in e.meta
             else f"[{e.dataset}] {e.task_name} ({e.meta['class']})" for e in embeddings]
    embeddings = np.array(embeddings)

    # === Plot comparison between taxonomical distance and task2vec distance ===
    # Compute all taxonomical distances
    tax_distance_matrix = np.zeros_like(distance_matrix)
    for i, e0 in enumerate(embeddings):
        for j, e1 in enumerate(embeddings):
            tax_distance_matrix[i, j] = taxonomy_distance(e0, e1)

    sns.set_style('whitegrid')
    plt.close('all')
    sns.set_style('whitegrid')
    fig = plt.figure(figsize=(3.3, 3.3))
    ax = fig.gca()
    # Plot how the average task2vec and taxonomical distance change as the neighborhood changes
    np.fill_diagonal(distance_matrix, np.inf)
    np.fill_diagonal(tax_distance_matrix, np.inf)
    plot_changing_k(ax, distance_matrix, embeddings, embeddings, label='Task2Vec distance')
    plot_changing_k(ax, tax_distance_matrix, embeddings, embeddings, label='Taxonomy distance')
    np.fill_diagonal(distance_matrix, 0)
    np.fill_diagonal(tax_distance_matrix, 0)
    ax.legend()
    # save figure
    fig.savefig('plots/embedding_distace_vs_taxonomy.pdf', bbox_inches='tight')

    # === Plot the clustered distance matrix ===
    np.fill_diagonal(distance_matrix, 0.)

    embeddings = embeddings[:50]
    names = names[:50]
    distance_matrix = distance_matrix[:50, :50]
    # sort distance matrix
    distance_matrix, embeddings, names, _ = sort_distance_matrix(distance_matrix, embeddings, names, method='complete')
    # draw figure
    plt.close('all')
    sns.set_style('white')
    draw_figure_to_plt(distance_matrix, embeddings, names, label_size=22)
    # save figure
    plt.savefig('plots/task2vec_distance_matrix.pdf', format=args.format, dpi=None, bbox_inches='tight')

if __name__ == '__main__':
    main()


import torch
import hashlib, os, pickle
import sys, time
from cassie.quaternion_function import *
import tty
import termios
import select
import numpy as np
from functools import partial
from rl.envs.wrappers import SymmetricEnv
from cassie import CassieEnv, CassiePlayground, CassieStandingEnv, CassieEnv_noaccel_footdist_omniscient, CassieEnv_noaccel_footdist

def isData():
    return select.select([sys.stdin], [], [], 0) == ([sys.stdin], [], [])

env = CassieEnv(state_est=True, dynamics_randomization=False, history=0)
env_fn = partial(CassieEnv, state_est=True, dynamics_randomization=False, history=0)
# env = CassieEnv_noaccel_footdist(state_est=True, dynamics_randomization=False, history=0)
# env_fn = partial(CassieEnv_noaccel_footdist, state_est=True, dynamics_randomization=False, history=0)

sym_env = SymmetricEnv(env_fn, mirrored_obs=env_fn().mirrored_obs, mirrored_act=[-5, -6, 7, 8, 9, -0.1, -1, 2, 3, 4])
# obs = env.get_full_state()
# print("obs len: ", len(obs))
# exit()

path = "./trained_models/nodelta_neutral_StateEst_symmetry_speed0-3_freq1-2/"
# path = "./logs/footdist/CassieNoaccelFootDist/noaccel_footdist_speedmatch_seed10/"
policy = torch.load(path + "actor.pt")
policy.eval()

old_settings = termios.tcgetattr(sys.stdin)

orient_add = 0

env.render()
render_state = True
try:
    tty.setcbreak(sys.stdin.fileno())

    state = env.reset_for_test()
    done = False
    timesteps = 0
    eval_reward = 0
    speed = 0.0

    while render_state:

        if isData():
            c = sys.stdin.read(1)
            if c == 'w':
                speed += 0.1
            elif c == 's':
                speed -= 0.1
            elif c == 'j':
                env.phase_add += .1
                print("Increasing frequency to: {:.1f}".format(env.phase_add))
            elif c == 'h':
                env.phase_add -= .1
                print("Decreasing frequency to: {:.1f}".format(env.phase_add))
            elif c == 'l':
                orient_add += .1
                print("Increasing orient_add to: ", orient_add)
            elif c == 'k':
                orient_add -= .1
                print("Decreasing orient_add to: ", orient_add)
            elif c == 'p':
                push = 100
                push_dir = 2
                force_arr = np.zeros(6)
                force_arr[push_dir] = push
                env.sim.apply_force(force_arr)

            env.update_speed(speed)
            print("speed: ", env.speed)

        if hasattr(env, 'simrate'):
            start = time.time()

        if (not env.vis.ispaused()):
            # Update Orientation
            quaternion = euler2quat(z=orient_add, y=0, x=0)
            iquaternion = inverse_quaternion(quaternion)

            if env.state_est:
                curr_orient = state[1:5]
                curr_transvel = state[15:18]
                # curr_orient = state[6:10]
                # curr_transvel = state[20:23]
            else:
                curr_orient = state[2:6]
                curr_transvel = state[20:23]

            new_orient = quaternion_product(iquaternion, curr_orient)

            if new_orient[0] < 0:
                new_orient = -new_orient

            new_translationalVelocity = rotate_by_quaternion(curr_transvel, iquaternion)

            if env.state_est:
                state[1:5] = torch.FloatTensor(new_orient)
                state[15:18] = torch.FloatTensor(new_translationalVelocity)
                # state[6:10] = torch.FloatTensor(new_orient)
                # state[20:23] = torch.FloatTensor(new_translationalVelocity)
                # state[0] = 1      # For use with StateEst. Replicate hack that height is always set to one on hardware.
            else:
                state[2:6] = torch.FloatTensor(new_orient)
                state[20:23] = torch.FloatTensor(new_translationalVelocity)

            state = torch.Tensor(state)
            # Calculate mirror state and mirror action
            with torch.no_grad():
                mirror_state = sym_env.mirror_clock_observation(state.unsqueeze(0), env.clock_inds)[0]
                # Mirror pelvis orientation and velocity
                # mir_quat = inverse_quaternion(mirror_state[1:5])
                # mir_quat[2] *= -1
                # mirror_state[1:5] = torch.Tensor(mir_quat)
                # mirror_state[16] *= -1      # y trans vel
                # mir_rot_vel = -mirror_state[18:21]
                # mir_rot_vel[1] *= -1
                # mirror_state[18:21] = mir_rot_vel
                # mirror_state[32] *= -1      # y trans accel
                mir_action = policy.forward(mirror_state, deterministic=True)
                mir_mir_action = sym_env.mirror_action(mir_action.unsqueeze(0)).detach().numpy()[0]
                action = policy.forward(state, deterministic=True).detach().numpy()
            # print("mirror action diff: ", np.linalg.norm(mir_mir_action - action))
            state, reward, done, _ = env.step(mir_mir_action)

            eval_reward += reward
            timesteps += 1

        render_state = env.render()
        if hasattr(env, 'simrate'):
            # assume 30hz (hack)
            end = time.time()
            delaytime = max(0, 1000 / 30000 - (end-start))
            time.sleep(delaytime)

    print("Eval reward: ", eval_reward)

finally:
    termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings)


import re
import torch
import pickle
import string
import numpy as np
from torch import nn
from torch.nn.utils.rnn import pack_padded_sequence
from pythainlp import word_tokenize
from pythainlp.util import emoji_to_thai

np.warnings.filterwarnings('ignore', category=np.VisibleDeprecationWarning)

class BILSTM_Model(torch.nn.Module) :
    def __init__(self, vocab_size, embedding_dim, hidden_dim, num_layers=1) :
        super().__init__()
        self.hidden_dim = hidden_dim
        self.dropout = nn.Dropout(0.3)
        self.embeddings = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim, num_layers, batch_first=True, bidirectional=True, dropout=0.2)
        self.linear1 = nn.Linear(hidden_dim*2, hidden_dim)
        self.linear2 = nn.Linear(hidden_dim, 4)

    def forward(self, x, s):
        x = self.embeddings(x)
        x = self.dropout(x)
        x_pack = pack_padded_sequence(x, s, batch_first=True, enforce_sorted=False)
        out_pack, (ht, ct) = self.lstm(x_pack)
        x = self.linear1(torch.cat((ht[-2,:,:], ht[-1,:,:]), dim = 1))
        return self.linear2(x)

def preprocess(text):
  def replace_rep(text):
    def _replace_rep(m):
      c,cc = m.groups()
      return f'{c}xxrep'
    re_rep = re.compile(r'(\S)(\1{2,})')
    return re_rep.sub(_replace_rep, text)

  def replace_url(text):
    URL_PATTERN = r"""(?i)\b((?:https?:(?:/{1,3}|[a-z0-9%])|[a-z0-9.\-]+[.](?:com|net|org|edu|gov|mil|aero|asia|biz|cat|coop|info|int|jobs|mobi|museum|name|post|pro|tel|travel|xxx|ac|ad|ae|af|ag|ai|al|am|an|ao|aq|ar|as|at|au|aw|ax|az|ba|bb|bd|be|bf|bg|bh|bi|bj|bm|bn|bo|br|bs|bt|bv|bw|by|bz|ca|cc|cd|cf|cg|ch|ci|ck|cl|cm|cn|co|cr|cs|cu|cv|cx|cy|cz|dd|de|dj|dk|dm|do|dz|ec|ee|eg|eh|er|es|et|eu|fi|fj|fk|fm|fo|fr|ga|gb|gd|ge|gf|gg|gh|gi|gl|gm|gn|gp|gq|gr|gs|gt|gu|gw|gy|hk|hm|hn|hr|ht|hu|id|ie|il|im|in|io|iq|ir|is|it|je|jm|jo|jp|ke|kg|kh|ki|km|kn|kp|kr|kw|ky|kz|la|lb|lc|li|lk|lr|ls|lt|lu|lv|ly|ma|mc|md|me|mg|mh|mk|ml|mm|mn|mo|mp|mq|mr|ms|mt|mu|mv|mw|mx|my|mz|na|nc|ne|nf|ng|ni|nl|no|np|nr|nu|nz|om|pa|pe|pf|pg|ph|pk|pl|pm|pn|pr|ps|pt|pw|py|qa|re|ro|rs|ru|rw|sa|sb|sc|sd|se|sg|sh|si|sj|Ja|sk|sl|sm|sn|so|sr|ss|st|su|sv|sx|sy|sz|tc|td|tf|tg|th|tj|tk|tl|tm|tn|to|tp|tr|tt|tv|tw|tz|ua|ug|uk|us|uy|uz|va|vc|ve|vg|vi|vn|vu|wf|ws|ye|yt|yu|za|zm|zw)/)(?:[^\s()<>{}\[\]]+|\([^\s()]*?\([^\s()]+\)[^\s()]*?\)|\([^\s]+?\))+(?:\([^\s()]*?\([^\s()]+\)[^\s()]*?\)|\([^\s]+?\)|[^\s`!()\[\]{};:'".,<>?«»“”‘’])|(?:(?<!@)[a-z0-9]+(?:[.\-][a-z0-9]+)*[.](?:com|net|org|edu|gov|mil|aero|asia|biz|cat|coop|info|int|jobs|mobi|museum|name|post|pro|tel|travel|xxx|ac|ad|ae|af|ag|ai|al|am|an|ao|aq|ar|as|at|au|aw|ax|az|ba|bb|bd|be|bf|bg|bh|bi|bj|bm|bn|bo|br|bs|bt|bv|bw|by|bz|ca|cc|cd|cf|cg|ch|ci|ck|cl|cm|cn|co|cr|cs|cu|cv|cx|cy|cz|dd|de|dj|dk|dm|do|dz|ec|ee|eg|eh|er|es|et|eu|fi|fj|fk|fm|fo|fr|ga|gb|gd|ge|gf|gg|gh|gi|gl|gm|gn|gp|gq|gr|gs|gt|gu|gw|gy|hk|hm|hn|hr|ht|hu|id|ie|il|im|in|io|iq|ir|is|it|je|jm|jo|jp|ke|kg|kh|ki|km|kn|kp|kr|kw|ky|kz|la|lb|lc|li|lk|lr|ls|lt|lu|lv|ly|ma|mc|md|me|mg|mh|mk|ml|mm|mn|mo|mp|mq|mr|ms|mt|mu|mv|mw|mx|my|mz|na|nc|ne|nf|ng|ni|nl|no|np|nr|nu|nz|om|pa|pe|pf|pg|ph|pk|pl|pm|pn|pr|ps|pt|pw|py|qa|re|ro|rs|ru|rw|sa|sb|sc|sd|se|sg|sh|si|sj|Ja|sk|sl|sm|sn|so|sr|ss|st|su|sv|sx|sy|sz|tc|td|tf|tg|th|tj|tk|tl|tm|tn|to|tp|tr|tt|tv|tw|tz|ua|ug|uk|us|uy|uz|va|vc|ve|vg|vi|vn|vu|wf|ws|ye|yt|yu|za|zm|zw)\b/?(?!@)))"""
    return re.sub(URL_PATTERN, 'xxurl', text)

  def replace_emoji(text):
    return emoji_to_thai(text, delimiters=(' ', ' '))

  def replace_punctuation(text):
    punctuation = string.punctuation.replace(".", "") # remove all punctuation except . because we want to keep infomation some word such as "อ.", "จ." and etc.
    punctuation_translator = str.maketrans('', '', punctuation)
    return text.translate(punctuation_translator)

  preprocess_text = text.lower().strip()
  preprocess_text = replace_url(preprocess_text)
  preprocess_text = replace_rep(preprocess_text)
  # preprocess_text = replace_emoji(preprocess_text)
  preprocess_text = replace_punctuation(preprocess_text)
  return preprocess_text

class MySentimentModel:

  def __init__(self) -> None:
      with open("./dumps/vocab2index.pkl", "rb") as f:
          self.vocab2index = pickle.load(f)
      with open("./dumps/id2label.pkl", "rb") as f:
          self.id2label = pickle.load(f)
      self.model = BILSTM_Model(len(self.vocab2index), 300, 256, 4)
      self.model.load_state_dict(torch.load('./weights/stacked_bilstm_model.pth', map_location=torch.device('cpu')))

  def encode_sentence(self, text, vocab2index, N=70):
      tokenized = word_tokenize(text)
      encoded = np.zeros(N, dtype=int)
      enc1 = np.array([vocab2index.get(word, vocab2index["UNK"]) for word in tokenized])
      length = min(N, len(enc1))
      encoded[:length] = enc1[:length]
      return encoded, length

  def predict(self, text) -> str:
      preprocess_text = preprocess(text)
      x, l = self.encode_sentence(preprocess_text,self.vocab2index )
      x, l = torch.Tensor(x).long(), torch.Tensor([l]).long()
      predict = self.model(x.unsqueeze(0), l)[0]
      prob = nn.functional.softmax(predict, dim=-1).detach().numpy().tolist()
      prediction = predict.detach().numpy()
      prediction_class = np.argmax(prediction)
      return self.id2label[prediction_class], prob


import spiceypy as spice
import numpy as np
import scipy.constants

import ale
from ale.base.type_sensor import Framer
from ale.transformation import FrameChain
from ale.rotation import TimeDependentRotation
from ale import util

class NaifSpice():
    def __enter__(self):
        """
        Called when the context is created. This is used
        to get the kernels furnished.
        """
        if self.kernels:
            [spice.furnsh(k) for k in self.kernels]
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        Called when the context goes out of scope. Once
        this is done, the object is out of scope and the
        kernels can be unloaded.
        """
        if self.kernels:
            [spice.unload(k) for k in self.kernels]

    @property
    def kernels(self):
        if not hasattr(self, '_kernels'):
            if 'kernels' in self._props.keys():
                try:
                    self._kernels = util.get_kernels_from_isis_pvl(self._props['kernels'])
                except Exception as e:
                    self._kernels =  self._props['kernels']
            else:
                if not ale.spice_root:
                    raise EnvironmentError(f'ale.spice_root is not set, cannot search for metakernels. ale.spice_root = "{ale.spice_root}"')

                search_results = util.get_metakernels(ale.spice_root, missions=self.short_mission_name, years=self.utc_start_time.year, versions='latest')

                if search_results['count'] == 0:
                    raise ValueError(f'Failed to find metakernels. mission: {self.short_mission_name}, year:{self.utc_start_time.year}, versions="latest" spice root = "{ale.spice_root}"')
                self._kernels = [search_results['data'][0]['path']]

        return self._kernels

    @property
    def light_time_correction(self):
        """
        Returns the type of light time correciton and abberation correction to
        use in NAIF calls. Expects ikid to be defined. This must be the integer
        Naif id code of the instrument.

        This searches for the value of the NAIF keyword INS<ikid>_LIGHTTIME_CORRECTION.
        If the keyword is not defined, then this defaults to light time
        correction and abberation correction (LT+S).

        Returns
        -------
        : str
          The light time and abberation correction string for use in NAIF calls.
          See https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/req/abcorr.html
          for the different options available.
        """
        try:
            return spice.gcpool('INS{}_LIGHTTIME_CORRECTION'.format(self.ikid), 0, 1)[0]
        except:
            return 'LT+S'

    @property
    def odtx(self):
        """
        Returns the x coefficient for the optical distortion model
        Expects ikid to be defined. This must be the integer Naif id code of the instrument

        Returns
        -------
        : list
          Optical distortion x coefficients
        """
        return spice.gdpool('INS{}_OD_T_X'.format(self.ikid),0, 10).tolist()

    @property
    def odty(self):
        """
        Returns the y coefficient for the optical distortion model.
        Expects ikid to be defined. This must be the integer Naif id code of the instrument

        Returns
        -------
        : list
          Optical distortion y coefficients
        """
        return spice.gdpool('INS{}_OD_T_Y'.format(self.ikid), 0, 10).tolist()

    @property
    def odtk(self):
        """
        The coefficients for the radial distortion model
        Expects ikid to be defined. This must be the integer Naif id code of the instrument

        Returns
        -------
        : list
          Radial distortion coefficients
        """
        return spice.gdpool('INS{}_OD_K'.format(self.ikid),0, 3).tolist()

    @property
    def ikid(self):
        """
        Returns the Naif ID code for the instrument
        Expects the instrument_id to be defined. This must be a string containing
        the short name of the instrument.

        Returns
        -------
        : int
          Naif ID used to for indentifying the instrument in Spice kernels
        """
        return spice.bods2c(self.instrument_id)

    @property
    def spacecraft_id(self):
        """
        Returns the Naif ID code for the spacecraft
        Expects the spacecraft_name to be defined. This must be a string containing
        the name of the spacecraft.

        Returns
        -------
        : int
          Naif ID code for the spacecraft
        """
        return spice.bods2c(self.spacecraft_name)

    @property
    def target_id(self):
        """
        Returns the Naif ID code for the target body
        Expects target_name to be defined. This must be a string containig the name
        of the target body.

        Returns
        -------
        : int
          Naif ID code for the target body
        """
        return spice.bods2c(self.target_name)

    @property
    def target_frame_id(self):
        """
        Returns the Naif ID code for the target reference frame
        Expects the target_id to be defined. This must be the integer Naif ID code
        for the target body.

        Returns
        -------
        : int
          Naif ID code for the target frame
        """
        frame_info = spice.cidfrm(self.target_id)
        return frame_info[0]

    @property
    def sensor_frame_id(self):
        """
        Returns the Naif ID code for the sensor reference frame
        Expects ikid to be defined. This must be the integer Naif id code of the instrument

        Returns
        -------
        : int
          Naif ID code for the sensor frame
        """
        return self.ikid

    @property
    def focal2pixel_lines(self):
        """
        Expects ikid to be defined. This must be the integer Naif id code of the instrument

        Returns
        -------
        : list<double>
          focal plane to detector lines
        """
        return list(spice.gdpool('INS{}_ITRANSL'.format(self.ikid), 0, 3))

    @property
    def focal2pixel_samples(self):
        """
        Expects ikid to be defined. This must be the integer Naif id code of the instrument

        Returns
        -------
        : list<double>
          focal plane to detector samples
        """
        return list(spice.gdpool('INS{}_ITRANSS'.format(self.ikid), 0, 3))

    @property
    def pixel2focal_x(self):
        """
        Expects ikid to be defined. This must be the integer Naif id code of the instrument

        Returns
        -------
        : list<double>
        detector to focal plane x
        """
        return list(spice.gdpool('INS{}_TRANSX'.format(self.ikid), 0, 3))

    @property
    def pixel2focal_y(self):
        """
        Expects ikid to be defined. This must be the integer Naif id code of the instrument

        Returns
        -------
        : list<double>
        detector to focal plane y
        """
        return list(spice.gdpool('INS{}_TRANSY'.format(self.ikid), 0, 3))

    @property
    def focal_length(self):
        """
        Returns the focal length of the sensor
        Expects ikid to be defined. This must be the integer Naif id code of the instrument

        Returns
        -------
        : float
          focal length
        """
        return float(spice.gdpool('INS{}_FOCAL_LENGTH'.format(self.ikid), 0, 1)[0])

    @property
    def pixel_size(self):
        """
        Expects ikid to be defined. This must be the integer Naif id code of the instrument

        Returns
        -------
        : float pixel size
        """
        return spice.gdpool('INS{}_PIXEL_SIZE'.format(self.ikid), 0, 1)[0] * 0.001

    @property
    def target_body_radii(self):
        """
        Returns a list containing the radii of the target body
        Expects target_name to be defined. This must be a string containing the name
        of the target body

        Returns
        -------
        : list<double>
          Radius of all three axis of the target body
        """
        rad = spice.bodvrd(self.target_name, 'RADII', 3)
        return rad[1]

    @property
    def reference_frame(self):
        """
        Returns a string containing the name of the target reference frame
        Expects target_name to be defined. This must be a string containing the name
        of the target body

        Returns
        -------
        : str
        String name of the target reference frame
        """
        try:
            return spice.cidfrm(spice.bodn2c(self.target_name))[1]
        except:
            return 'IAU_{}'.format(self.target_name)

    @property
    def sun_position(self):
        """
        Returns a tuple with information detailing the sun position at the time
        of the image. Expects center_ephemeris_time to be defined. This must be
        a floating point number containing the average of the start and end ephemeris time.
        Expects reference frame to be defined. This must be a sring containing the name of
        the target reference frame. Expects target_name to be defined. This must be
        a string containing the name of the target body.

        Returns
        -------
        : (sun_positions, sun_velocities)
          a tuple containing a list of sun positions, a list of sun velocities
        """
        times = [self.center_ephemeris_time]
        positions = []
        velocities = []

        for time in times:
            sun_state, _ = spice.spkezr("SUN",
                                         time,
                                         self.reference_frame,
                                         'LT+S',
                                         self.target_name)
            positions.append(sun_state[:3])
            velocities.append(sun_state[3:6])
        positions = 1000 * np.asarray(positions)
        velocities = 1000 * np.asarray(velocities)

        return positions, velocities, times

    @property
    def sensor_position(self):
        """
        Returns a tuple with information detailing the position of the sensor at the time
        of the image. Expects ephemeris_time to be defined. This must be a floating point number
        containing the ephemeris time. Expects spacecraft_name to be defined. This must be a
        string containing the name of the spacecraft containing the sensor. Expects
        reference_frame to be defined. This must be a sring containing the name of
        the target reference frame. Expects target_name to be defined. This must be
        a string containing the name of the target body.

        Returns
        -------
        : (positions, velocities, times)
          a tuple containing a list of positions, a list of velocities, and a list of times
        """
        if not hasattr(self, '_position'):
            ephem = self.ephemeris_time
            pos = []
            vel = []

            target = self.spacecraft_name
            observer = self.target_name
            # Check for ISIS flag to fix target and observer swapping
            if self.swap_observer_target:
                target = self.target_name
                observer = self.spacecraft_name

            for time in ephem:
                # spkezr returns a vector from the observer's location to the aberration-corrected
                # location of the target. For more information, see:
                # https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/FORTRAN/spicelib/spkezr.html
                if self.correct_lt_to_surface and self.light_time_correction.upper() == 'LT+S':
                    obs_tar_state, obs_tar_lt = spice.spkezr(target,
                                                             time,
                                                             'J2000',
                                                             self.light_time_correction,
                                                             observer)
                    # ssb to spacecraft
                    ssb_obs_state, ssb_obs_lt = spice.spkezr(observer,
                                                             time,
                                                             'J2000',
                                                             'NONE',
                                                             'SSB')

                    radius_lt = (self.target_body_radii[2] + self.target_body_radii[0]) / 2 / (scipy.constants.c/1000.0)
                    adjusted_time = time - obs_tar_lt + radius_lt
                    ssb_tar_state, ssb_tar_lt = spice.spkezr(target,
                                                             adjusted_time,
                                                             'J2000',
                                                             'NONE',
                                                             'SSB')
                    state = ssb_tar_state - ssb_obs_state
                    matrix = spice.sxform("J2000", self.reference_frame, time)
                    state = spice.mxvg(matrix, state, 6, 6);
                else:
                    state, _ = spice.spkezr(target,
                                            time,
                                            self.reference_frame,
                                            self.light_time_correction,
                                            observer)

                if self.swap_observer_target:
                    pos.append(-state[:3])
                    vel.append(-state[3:])
                else:
                    pos.append(state[:3])
                    vel.append(state[3:])

            # By default, SPICE works in km, so convert to m
            self._position = [p * 1000 for p in pos]
            self._velocity = [v * 1000 for v in vel]
        return self._position, self._velocity, self.ephemeris_time

    @property
    def frame_chain(self):
        if not hasattr(self, '_frame_chain'):
            nadir = self._props.get('nadir', False)
            self._frame_chain = FrameChain.from_spice(sensor_frame=self.sensor_frame_id,
                                                      target_frame=self.target_frame_id,
                                                      center_ephemeris_time=self.center_ephemeris_time,
                                                      ephemeris_times=self.ephemeris_time,
                                                      nadir=nadir)

            if nadir:
                # Logic for nadir calculation was taken from ISIS3
                #  SpiceRotation::setEphemerisTimeNadir
                rotation = self._frame_chain.compute_rotation(self.target_frame_id, 1)
                p_vec, v_vec, times = self.sensor_position
                rotated_positions = rotation.apply_at(p_vec, times)
                rotated_velocities = rotation.rotate_velocity_at(p_vec, v_vec, times)

                p_vec = rotated_positions
                v_vec = rotated_velocities

                velocity_axis = 2
                # Get the default line translation with no potential flipping
                # from the driver
                trans_x = np.array(list(spice.gdpool('INS{}_ITRANSL'.format(self.ikid), 0, 3)))

                if (trans_x[0] < trans_x[1]):
                    velocity_axis = 1

                quats = [spice.m2q(spice.twovec(-p_vec[i], 3, v_vec[i], velocity_axis)) for i, time in enumerate(times)]
                quats = np.array(quats)[:,[1,2,3,0]]

                rotation = TimeDependentRotation(quats, times, 1, self.sensor_frame_id)
                self._frame_chain.add_edge(rotation)

        return self._frame_chain

    @property
    def sensor_orientation(self):
        """
        Returns quaternions describing the sensor orientation. Expects ephemeris_time
        to be defined. This must be a floating point number containing the
        ephemeris time. Expects instrument_id to be defined. This must be a string
        containing the short name of the instrument. Expects reference frame to be defined.
        This must be a string containing the name of the target reference frame.

        Returns
        -------
        : list
          Quaternions describing the orientation of the sensor
        """
        if not hasattr(self, '_orientation'):
            self._orientation = self.frame_chain.compute_rotation(self.sensor_frame_id, self.target_frame_id).quats
        return self._orientation.tolist()

    @property
    def ephemeris_start_time(self):
        """
        Returns the starting ephemeris time of the image. Expects spacecraft_id to
        be defined. This must be the integer Naif Id code for the spacecraft. Expects
        spacecraft_clock_start_count to be defined. This must be a string
        containing the start clock count of the spacecraft

        Returns
        -------
        : double
          Starting ephemeris time of the image
        """
        return spice.scs2e(self.spacecraft_id, self.spacecraft_clock_start_count)

    @property
    def ephemeris_stop_time(self):
        """
        Returns the ephemeris stop time of the image. Expects spacecraft_id to
        be defined. This must be the integer Naif Id code for the spacecraft.
        Expects spacecraft_clock_stop_count to be defined. This must be a string
        containing the stop clock count of the spacecraft

        Returns
        -------
        : double
          Ephemeris stop time of the image
        """
        return spice.scs2e(self.spacecraft_id, self.spacecraft_clock_stop_count)

    @property
    def detector_center_sample(self):
        """
        Returns the center detector sample. Expects ikid to be defined. This should
        be an integer containing the Naif Id code of the instrument.

        Returns
        -------
        : float
          Detector sample of the principal point
        """
        return float(spice.gdpool('INS{}_BORESIGHT_SAMPLE'.format(self.ikid), 0, 1)[0])

    @property
    def detector_center_line(self):
        """
        Returns the center detector line. Expects ikid to be defined. This should
        be an integer containing the Naif Id code of the instrument.

        Returns
        -------
        : float
          Detector line of the principal point
        """
        return float(spice.gdpool('INS{}_BORESIGHT_LINE'.format(self.ikid), 0, 1)[0])

    @property
    def swap_observer_target(self):
        """
        Returns if the observer and target should be swapped when determining the
        sensor state relative to the target. This is defined by a keyword in
        ISIS IAKs. If the keyword is not defined in any loaded kernels then False
        is returned.

        Expects ikid to be defined. This should be an integer containing the
        Naif Id code of the instrument.
        """
        try:
            swap = spice.gcpool('INS{}_SWAP_OBSERVER_TARGET'.format(self.ikid), 0, 1)[0]
            return swap.upper() == "TRUE"
        except:
            return False

    @property
    def correct_lt_to_surface(self):
        """
        Returns if light time correction should be made to the surface instead of
        to the center of the body. This is defined by a keyword in ISIS IAKs.
        If the keyword is not defined in any loaded kernels then False is returned.

        Expects ikid to be defined. This should be an integer containing the
        Naif Id code of the instrument.
        """
        try:
            surface_correct = spice.gcpool('INS{}_LT_SURFACE_CORRECT'.format(self.ikid), 0, 1)[0]
            return surface_correct.upper() == "TRUE"
        except:
            return False

    @property
    def naif_keywords(self):
        """
        Returns
        -------
        : dict
          Dictionary of keywords and values that ISIS creates and attaches to the label
        """
        if not hasattr(self, "_naif_keywords"):
            self._naif_keywords = dict()

            self._naif_keywords['BODY{}_RADII'.format(self.target_id)] = self.target_body_radii
            self._naif_keywords['BODY_FRAME_CODE'] = self.target_frame_id
            self._naif_keywords['BODY_CODE'] = self.target_id

            self._naif_keywords = {**self._naif_keywords, **util.query_kernel_pool(f"*{self.ikid}*"),  **util.query_kernel_pool(f"*{self.target_id}*")}

            try:
                self._naif_keywords = {**self._naif_keywords, **util.query_kernel_pool(f"*{self.fikid}*")}
            except AttributeError as error:
                pass

        return self._naif_keywords


#!/usr/bin/env python
#
# This file contains a reference implementation of the following paper:
#   On Learning Associations of Faces and Voices
#   Changil Kim, Hijung Valentina Shin, Tae-Hyun Oh, Alexandre Kaspar, Mohamed Elgharib, Wojciech Matusik
#   ACCV 2018
# Please cite the above paper if you use this software.
# More information including a pre-trained model can be found in the project website:
#   http://facevoice.csail.mit.edu/
#
# Usage:
#   voice-to-face:  facevoice.py v2f -c <checkpoint> --voice <voice-file> --face0 <face-file> --face1 <face-file>
#   face-to-voice:  facevoice.py f2v -c <checkpoint> --face <face-file> --voice0 <voice-file> --voice1 <voice-file>
#
# Author: Changil Kim <changil@csail.mit.edu>
#

from __future__ import unicode_literals
from functools import reduce
import tensorflow as tf
import numpy as np
import warnings
import argparse
import skimage.io
import skimage.transform
import skimage
import scipy.io.wavfile

def conv(x, kh, kw, co, sh, sw, name, relu=True, padding='SAME'):
    with tf.variable_scope(name):
        ci = x.get_shape()[-1]
        weights = tf.get_variable('weights', shape=[kh, kw, ci, co])
        biases = tf.get_variable('biases', shape=[co])
        x = tf.nn.conv2d(x, weights, [1, sh, sw, 1], padding=padding)
        x = tf.nn.bias_add(x, biases)
        if relu:
            x = tf.nn.relu(x)
        return x

def fc(x, co, name, relu=True):
    with tf.variable_scope(name):
        # flatten the input tensor
        ci = reduce(lambda a, b: a*b, x.get_shape()[1:].as_list())
        x = tf.reshape(x, [-1, ci])
        weights = tf.get_variable('weights', shape=[ci, co])
        biases = tf.get_variable('biases', shape=[co])
        if relu:
            x = tf.nn.relu_layer(x, weights, biases)
        else:
            x = tf.nn.xw_plus_b(x, weights, biases)
        return x

def bn(x, name, relu=True):
    # for inference only
    with tf.variable_scope(name):
        ci = x.get_shape()[-1]
        scale = tf.get_variable('scale', shape=[ci])
        offset = tf.get_variable('offset', shape=[ci])
        mean = tf.get_variable('mean', shape=[ci])
        variance = tf.get_variable('variance', shape=[ci])
        x = tf.nn.batch_normalization(x, mean=mean, variance=variance, offset=offset, scale=scale, variance_epsilon=1e-5, name=name)
        if relu:
            x = tf.nn.relu(x)
        return x

def maxpool(x, kh, kw, sh, sw, name, padding='SAME'):
    return tf.nn.max_pool(x, ksize=[1, kh, kw, 1], strides=[1, sh, sw, 1], padding=padding, name=name)

def avgpool(x, kh, kw, sh, sw, name, padding='SAME'):
    return tf.nn.avg_pool(x, ksize=[1, kh, kw, 1], strides=[1, sh, sw, 1], padding=padding, name=name)

def padding(x, h, w, name):
    # for torch-style padding
    pads = np.zeros([4, 2], dtype=np.int32)
    pads[1, :] = h
    pads[2, :] = w
    return tf.pad(x, tf.constant(pads), name=name)

def facenet(x):
    # vgg16 up to conv5_3
    x = conv(x, 3, 3, 64, 1, 1, name='conv1_1')
    x = conv(x, 3, 3, 64, 1, 1, name='conv1_2')
    x = maxpool(x, 2, 2, 2, 2, name='pool1')
    x = conv(x, 3, 3, 128, 1, 1, name='conv2_1')
    x = conv(x, 3, 3, 128, 1, 1, name='conv2_2')
    x = maxpool(x, 2, 2, 2, 2, name='pool2')
    x = conv(x, 3, 3, 256, 1, 1, name='conv3_1')
    x = conv(x, 3, 3, 256, 1, 1, name='conv3_2')
    x = conv(x, 3, 3, 256, 1, 1, name='conv3_3')
    x = maxpool(x, 2, 2, 2, 2, name='pool3')
    x = conv(x, 3, 3, 512, 1, 1, name='conv4_1')
    x = conv(x, 3, 3, 512, 1, 1, name='conv4_2')
    x = conv(x, 3, 3, 512, 1, 1, name='conv4_3')
    x = maxpool(x, 2, 2, 2, 2, name='pool4')
    x = conv(x, 3, 3, 512, 1, 1, name='conv5_1')
    x = conv(x, 3, 3, 512, 1, 1, name='conv5_2')
    x = conv(x, 3, 3, 512, 1, 1, name='conv5_3')
    # global pooling
    x = avgpool(x, 14, 14, 1, 1, padding='VALID', name='globalpool')
    # fc layers
    x = fc(x, 128, name='fc1')
    x = fc(x, 128, relu=False, name='fc2')
    return x

def voicenet(x):
    # soundnet up to conv6
    x = padding(x, 32, 0, name='pad1')
    x = conv(x, 64, 1, 16, 2, 1, padding='VALID', relu=False, name='conv1')
    x = bn(x, relu=True, name='bn1')
    x = maxpool(x, 8, 1, 8, 1, padding='VALID', name='pool1')
    x = padding(x, 16, 0, name='pad2')
    x = conv(x, 32, 1, 32, 2, 1, padding='VALID', relu=False, name='conv2')
    x = bn(x, relu=True, name='bn2')
    x = maxpool(x, 8, 1, 8, 1, padding='VALID', name='pool2')
    x = padding(x, 8, 0, name='pad3')
    x = conv(x, 16, 1, 64, 2, 1, padding='VALID', relu=False, name='conv3')
    x = bn(x, relu=True, name='bn3')
    x = padding(x, 4, 0, name='pad4')
    x = conv(x, 8, 1, 128, 2, 1, padding='VALID', relu=False, name='conv4')
    x = bn(x, relu=True, name='bn4')
    x = padding(x, 2, 0, name='pad5')
    x = conv(x, 4, 1, 256, 2, 1, padding='VALID', relu=False, name='conv5')
    x = bn(x, relu=True, name='bn5')
    x = maxpool(x, 4, 1, 4, 1, padding='VALID', name='pool5')
    x = padding(x, 2, 0, name='pad6')
    x = conv(x, 4, 1, 512, 2, 1, padding='VALID', relu=False, name='conv6')
    # global pooling
    x = avgpool(x, 14, 1, 1, 1, padding='VALID', name='globalpool')
    # fc layers
    x = fc(x, 128, name='fc1')
    x = fc(x, 128, relu=False, name='fc2')
    return x

def voice2face(voice, face0, face1):
    with tf.variable_scope('voice'):
        voicerep = voicenet(voice)
    with tf.variable_scope('face'):
        facerep0 = facenet(face0)
    with tf.variable_scope('face', reuse=True):
        facerep1 = facenet(face1)

    # pick the closer face
    return which(voicerep, facerep0, facerep1)

def face2voice(face, voice0, voice1):
    with tf.variable_scope('face'):
        facerep = facenet(face)
    with tf.variable_scope('voice'):
        voicerep0 = voicenet(voice0)
    with tf.variable_scope('voice', reuse=True):
        voicerep1 = voicenet(voice1)

    # pick the closer voice
    return which(facerep, voicerep0, voicerep1)

def which(repref, rep0, rep1):
    # pick the one closer to repref between rep0 and rep1
    dist_sq = lambda a, b: tf.reduce_sum(tf.square(a - b), axis=1)
    d0 = dist_sq(repref, rep0)
    d1 = dist_sq(repref, rep1)
    batch_size = tf.shape(d0)[0]
    verdict = tf.where(d0 < d1, tf.fill(dims=[batch_size], value=0), tf.fill(dims=[batch_size], value=1))
    return verdict

def load_wave(wave_file):
    rate, data = scipy.io.wavfile.read(wave_file)
    if rate != 22050:
        raise RuntimeError('input wav must be sampled at 22,050 Hz')
    if data.ndim > 1:
        # take the left channel
        data = data[:, 0]
    if data.shape[0] < 220500:
        # make the wav at least 10-second long
        data = np.tile(data, (220500 + data.shape[0] - 1) // data.shape[0])
    # take the first 10 seconds
    data = np.reshape(data[:220500], [-1, 1, 1]).astype(np.float32)
    return data

def load_image(image_file):
    image = skimage.io.imread(image_file)
    if image.ndim != 3 or image.shape[2] != 3:
        raise RuntimeError('input image must be in color')
    # resize
    image = skimage.transform.resize(image, [256, 256], mode='reflect', anti_aliasing=True)
    with warnings.catch_warnings():
        warnings.simplefilter('ignore')
        image = skimage.img_as_ubyte(image)
    # crop
    image = image[16:-16, 16:-16]
    # subtract the mean
    image = image.astype(np.float32) - np.array([124., 117., 104.], dtype=np.float32)
    # RGB to BGR
    image = image[:, :, ::-1]
    return image

def testv2f(opt):
    # load files
    voice = np.expand_dims(load_wave(opt.voice), axis=0)
    face0 = np.expand_dims(load_image(opt.face0), axis=0)
    face1 = np.expand_dims(load_image(opt.face1), axis=0)

    # create placeholders for the input
    voicenode = tf.placeholder(tf.float32, shape=(None, 220500, 1, 1), name='voicenode')
    facenode0 = tf.placeholder(tf.float32, shape=(None, 224, 224, 3), name='facenode0')
    facenode1 = tf.placeholder(tf.float32, shape=(None, 224, 224, 3), name='facenode1')

    # set up network
    verdictnode = voice2face(voicenode, facenode0, facenode1)

    saver = tf.train.Saver(save_relative_paths=True)

    with tf.Session() as sess:
        # load checkpoint
        saver.restore(sess, tf.train.latest_checkpoint(opt.checkpoint_dir))

        # test
        verdict = sess.run(verdictnode, feed_dict={voicenode: voice, facenode0: face0, facenode1: face1})

        print('predicted matching face: face{} ({})'.format(verdict[0], [opt.face0, opt.face1][verdict[0]]))

def testf2v(opt):
    # load files
    face = np.expand_dims(load_image(opt.face), axis=0)
    voice0 = np.expand_dims(load_wave(opt.voice0), axis=0)
    voice1 = np.expand_dims(load_wave(opt.voice1), axis=0)

    # create placeholders for the input
    facenode = tf.placeholder(tf.float32, shape=(None, 224, 224, 3), name='facenode')
    voicenode0 = tf.placeholder(tf.float32, shape=(None, 220500, 1, 1), name='voicenode0')
    voicenode1 = tf.placeholder(tf.float32, shape=(None, 220500, 1, 1), name='voicenode1')

    # set up network
    verdictnode = face2voice(facenode, voicenode0, voicenode1)

    saver = tf.train.Saver(save_relative_paths=True)

    with tf.Session() as sess:
        # load checkpoint
        saver.restore(sess, tf.train.latest_checkpoint(opt.checkpoint_dir))

        # test
        verdict = sess.run(verdictnode, feed_dict={facenode: face, voicenode0: voice0, voicenode1: voice1})

        print('predicted matching voice: voice{} ({})'.format(verdict[0], [opt.voice0, opt.voice1][verdict[0]]))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers()
    # v2f
    parserv2f = subparsers.add_parser('v2f')
    parserv2f.add_argument('-c', '--checkpoint-dir', required=True)
    parserv2f.add_argument('--voice', required=True)
    parserv2f.add_argument('--face0', required=True)
    parserv2f.add_argument('--face1', required=True)
    parserv2f.set_defaults(func=testv2f)
    # f2v
    parserf2v = subparsers.add_parser('f2v')
    parserf2v.add_argument('-c', '--checkpoint-dir', required=True)
    parserf2v.add_argument('--face', required=True)
    parserf2v.add_argument('--voice0', required=True)
    parserf2v.add_argument('--voice1', required=True)
    parserf2v.set_defaults(func=testf2v)

    opt = parser.parse_args()
    opt.func(opt)


# This code is based on: https://github.com/nutonomy/second.pytorch.git
#
# MIT License
# Copyright (c) 2018
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn

class Scalar(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer('total', torch.FloatTensor([0.0]))
        self.register_buffer('count', torch.FloatTensor([0.0]))

    def forward(self, scalar):
        if not scalar.eq(0.0):
            self.count += 1
            self.total += scalar.data.float()
        return self.value.cpu()

    @property
    def value(self):
        return self.total / self.count

    def clear(self):
        self.total.zero_()
        self.count.zero_()

class Accuracy(nn.Module):
    def __init__(self,
                 dim=1,
                 ignore_idx=-1,
                 threshold=0.5,
                 encode_background_as_zeros=True):
        super().__init__()
        self.register_buffer('total', torch.FloatTensor([0.0]))
        self.register_buffer('count', torch.FloatTensor([0.0]))
        self._ignore_idx = ignore_idx
        self._dim = dim
        self._threshold = threshold
        self._encode_background_as_zeros = encode_background_as_zeros

    def forward(self, labels, preds, weights=None):
        # labels: [N, ...]
        # preds: [N, C, ...]
        if self._encode_background_as_zeros:
            scores = torch.sigmoid(preds)
            labels_pred = torch.max(preds, dim=self._dim)[1] + 1
            pred_labels = torch.where((scores > self._threshold).any(self._dim),
                                      labels_pred,
                                      torch.tensor(0).type_as(labels_pred))
        else:
            pred_labels = torch.max(preds, dim=self._dim)[1]
        N, *Ds = labels.shape
        labels = labels.view(N, int(np.prod(Ds)))
        pred_labels = pred_labels.view(N, int(np.prod(Ds)))
        if weights is None:
            weights = (labels != self._ignore_idx).float()
        else:
            weights = weights.float()

        num_examples = torch.sum(weights)
        num_examples = torch.clamp(num_examples, min=1.0).float()
        total = torch.sum((pred_labels == labels.long()).float())
        self.count += num_examples
        self.total += total
        return self.value.cpu()
        # return (total /  num_examples.data).cpu()
    @property
    def value(self):
        return self.total / self.count

    def clear(self):
        self.total.zero_()
        self.count.zero_()

class Precision(nn.Module):
    def __init__(self, dim=1, ignore_idx=-1, threshold=0.5):
        super().__init__()
        self.register_buffer('total', torch.FloatTensor([0.0]))
        self.register_buffer('count', torch.FloatTensor([0.0]))
        self._ignore_idx = ignore_idx
        self._dim = dim
        self._threshold = threshold

    def forward(self, labels, preds, weights=None):
        # labels: [N, ...]
        # preds: [N, C, ...]
        if preds.shape[self._dim] == 1:  # BCE
            pred_labels = (torch.sigmoid(preds) >
                           self._threshold).long().squeeze(self._dim)
        else:
            assert preds.shape[
                self._dim] == 2, "precision only support 2 class"
            pred_labels = torch.max(preds, dim=self._dim)[1]
        N, *Ds = labels.shape
        labels = labels.view(N, int(np.prod(Ds)))
        pred_labels = pred_labels.view(N, int(np.prod(Ds)))
        if weights is None:
            weights = (labels != self._ignore_idx).float()
        else:
            weights = weights.float()

        pred_trues = pred_labels > 0
        pred_falses = pred_labels == 0
        trues = labels > 0
        falses = labels == 0
        true_positives = (weights * (trues & pred_trues).float()).sum()
        true_negatives = (weights * (falses & pred_falses).float()).sum()
        false_positives = (weights * (falses & pred_trues).float()).sum()
        false_negatives = (weights * (trues & pred_falses).float()).sum()
        count = true_positives + false_positives
        # print(count, true_positives)
        if count > 0:
            self.count += count
            self.total += true_positives
        return self.value.cpu()
        # return (total /  num_examples.data).cpu()
    @property
    def value(self):
        return self.total / self.count
    def clear(self):
        self.total.zero_()
        self.count.zero_()

class Recall(nn.Module):
    def __init__(self, dim=1, ignore_idx=-1, threshold=0.5):
        super().__init__()
        self.register_buffer('total', torch.FloatTensor([0.0]))
        self.register_buffer('count', torch.FloatTensor([0.0]))
        self._ignore_idx = ignore_idx
        self._dim = dim
        self._threshold = threshold

    def forward(self, labels, preds, weights=None):
        # labels: [N, ...]
        # preds: [N, C, ...]
        if preds.shape[self._dim] == 1:  # BCE
            pred_labels = (torch.sigmoid(preds) >
                           self._threshold).long().squeeze(self._dim)
        else:
            assert preds.shape[
                self._dim] == 2, "precision only support 2 class"
            pred_labels = torch.max(preds, dim=self._dim)[1]
        N, *Ds = labels.shape
        labels = labels.view(N, int(np.prod(Ds)))
        pred_labels = pred_labels.view(N, int(np.prod(Ds)))
        if weights is None:
            weights = (labels != self._ignore_idx).float()
        else:
            weights = weights.float()
        pred_trues = pred_labels == 1
        pred_falses = pred_labels == 0
        trues = labels == 1
        falses = labels == 0
        true_positives = (weights * (trues & pred_trues).float()).sum()
        true_negatives = (weights * (falses & pred_falses).float()).sum()
        false_positives = (weights * (falses & pred_trues).float()).sum()
        false_negatives = (weights * (trues & pred_falses).float()).sum()
        count = true_positives + false_negatives
        if count > 0:
            self.count += count
            self.total += true_positives
        return self.value.cpu()
        # return (total /  num_examples.data).cpu()
    @property
    def value(self):
        return self.total / self.count
    def clear(self):
        self.total.zero_()
        self.count.zero_()

def _calc_binary_metrics(labels,
                         scores,
                         weights=None,
                         ignore_idx=-1,
                         threshold=0.5):

    pred_labels = (scores > threshold).long()
    N, *Ds = labels.shape
    labels = labels.view(N, int(np.prod(Ds)))
    pred_labels = pred_labels.view(N, int(np.prod(Ds)))
    pred_trues = pred_labels > 0
    pred_falses = pred_labels == 0
    trues = labels > 0
    falses = labels == 0
    true_positives = (weights * (trues & pred_trues).float()).sum()
    true_negatives = (weights * (falses & pred_falses).float()).sum()
    false_positives = (weights * (falses & pred_trues).float()).sum()
    false_negatives = (weights * (trues & pred_falses).float()).sum()
    return true_positives, true_negatives, false_positives, false_negatives

class PrecisionRecall(nn.Module):
    def __init__(self,
                 dim=1,
                 ignore_idx=-1,
                 thresholds=0.5,
                 use_sigmoid_score=False,
                 encode_background_as_zeros=True):
        super().__init__()
        if not isinstance(thresholds, (list, tuple)):
            thresholds = [thresholds]

        self.register_buffer('prec_total',
                             torch.FloatTensor(len(thresholds)).zero_())
        self.register_buffer('prec_count',
                             torch.FloatTensor(len(thresholds)).zero_())
        self.register_buffer('rec_total',
                             torch.FloatTensor(len(thresholds)).zero_())
        self.register_buffer('rec_count',
                             torch.FloatTensor(len(thresholds)).zero_())

        self._ignore_idx = ignore_idx
        self._dim = dim
        self._thresholds = thresholds
        self._use_sigmoid_score = use_sigmoid_score
        self._encode_background_as_zeros = encode_background_as_zeros

    def forward(self, labels, preds, weights=None):
        # labels: [N, ...]
        # preds: [N, ..., C]
        if self._encode_background_as_zeros:
            # this don't support softmax
            assert self._use_sigmoid_score is True
            total_scores = torch.sigmoid(preds)
            # scores, label_preds = torch.max(total_scores, dim=1)
        else:
            if self._use_sigmoid_score:
                total_scores = torch.sigmoid(preds)[..., 1:]
            else:
                total_scores = F.softmax(preds, dim=-1)[..., 1:]
        """
        if preds.shape[self._dim] == 1:  # BCE
            scores = torch.sigmoid(preds)
        else:
            # assert preds.shape[
            #     self._dim] == 2, "precision only support 2 class"
            # TODO: add support for [N, C, ...] format.
            # TODO: add multiclass support
            if self._use_sigmoid_score:
                scores = torch.sigmoid(preds)[:, ..., 1:].sum(-1)
            else:
                scores = F.softmax(preds, dim=self._dim)[:, ..., 1:].sum(-1)
        """
        scores = torch.max(total_scores, dim=-1)[0]
        if weights is None:
            weights = (labels != self._ignore_idx).float()
        else:
            weights = weights.float()
        for i, thresh in enumerate(self._thresholds):
            tp, tn, fp, fn = _calc_binary_metrics(labels, scores, weights,
                                                  self._ignore_idx, thresh)
            rec_count = tp + fn
            prec_count = tp + fp
            if rec_count > 0:
                self.rec_count[i] += rec_count
                self.rec_total[i] += tp
            if prec_count > 0:
                self.prec_count[i] += prec_count
                self.prec_total[i] += tp

        return self.value
        # return (total /  num_examples.data).cpu()
    @property
    def value(self):
        prec_count = torch.clamp(self.prec_count, min=1.0)
        rec_count = torch.clamp(self.rec_count, min=1.0)
        return ((self.prec_total / prec_count).cpu(),
                (self.rec_total / rec_count).cpu())

    @property
    def thresholds(self):
        return self._thresholds

    def clear(self):
        self.rec_count.zero_()
        self.prec_count.zero_()
        self.prec_total.zero_()
        self.rec_total.zero_()


#  Copyright 2013-2021 The Salish Sea MEOPAR contributors
#  and The University of British Columbia
#
#  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
#
#     https://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.
"""Produce a figure that shows water levels at a tide gauge station
calculated by the VHFR FVCOM and SalishSeaCast NEMO models,
and predicted and observed water levels from the CHS
:kbd:`https://ws-shc.qc.dfo-mpo.gc.ca/` water levels web service.

Testing notebook for this module is
https://nbviewer.jupyter.org/github/SalishSeaCast/SalishSeaNowcast/blob/main/notebooks/figures/fvcom/publish/TestTideStnWaterLevel.ipynb

Development notebook for this module is
https://nbviewer.jupyter.org/github/SalishSeaCast/SalishSeaNowcast/blob/main/notebooks/figures/fvcom/publish/DevelopTideStnWaterLevel.ipynb
"""
from contextlib import suppress
from datetime import timedelta
from pathlib import Path
from types import SimpleNamespace

import arrow
import matplotlib.dates
import matplotlib.pyplot as plt
import numpy
import requests
import xarray
from pandas.plotting import register_matplotlib_converters
from salishsea_tools import data_tools
from salishsea_tools.places import PLACES

import nowcast.figures.website_theme
from nowcast.figures import shared

def make_figure(
    place,
    fvcom_ssh_datasets,
    nemo_ssh_dataset_url_tmpl,
    figsize=(16, 9),
    theme=nowcast.figures.website_theme,
):
    """Plot water levels calculated by the VHFR FVCOM and SalishSeaCast NEMO
    models, and predicted and observed water levels from the CHS
    :kbd:`https://ws-shc.qc.dfo-mpo.gc.ca/` water levels web service for the
    tide gauge station at :kbd:`place`.

    :arg str place: Tide gauge station name;
                    must be a key in :py:obj:`salishsea_tools.places.PLACES`.

    :arg fvcom_ssh_datasets: Dictionary of VHFR FVCOM model tide gauge station
                             sea surface height time series of
                             py:class:`xarray.Dataset` objects keyed by model
                             model configuration (:kbd:`x2`, :kbd:`r12`).
    :type fvcom_ssh_datasets: :py:class:`dict`

    :arg str nemo_ssh_dataset_url_tmpl: ERDDAP URL template for SalishSeaCast
                                        NEMO model tide gauge station
                                        sea surface height time series dataset.

    :arg 2-tuple figsize: Figure size (width, height) in inches.

    :arg theme: Module-like object that defines the style elements for the
                figure. See :py:mod:`nowcast.figures.website_theme` for an
                example.

    :returns: :py:class:`matplotlib.figure.Figure`
    """
    plot_data = _prep_plot_data(place, fvcom_ssh_datasets, nemo_ssh_dataset_url_tmpl)
    fig, (ax_ssh, ax_res) = _prep_fig_axes(figsize, theme)
    _plot_water_level_time_series(ax_ssh, plot_data, theme)
    _water_level_time_series_labels(ax_ssh, plot_data, place, theme)
    _plot_residual_time_series(ax_res, plot_data, theme)
    _residual_time_series_labels(ax_res, plot_data, theme)
    return fig

def _prep_plot_data(place, fvcom_ssh_datasets, nemo_ssh_dataset_url_tmpl):
    # FVCOM sea surface height dataset(s)
    fvcom_sshs = {}
    for model_config, fvcom_ssh_dataset in fvcom_ssh_datasets.items():
        fvcom_sshs[model_config] = fvcom_ssh_dataset.zeta.isel(
            station=[
                name.decode().strip().split(maxsplit=1)[1]
                for name in fvcom_ssh_dataset.name_station.values
            ].index(place)
        )
        # Drop repeated times
        _, index = numpy.unique(fvcom_sshs[model_config].time.values, return_index=True)
        fvcom_sshs[model_config] = fvcom_sshs[model_config].isel(time=index)
    fvcom_period = slice(
        str(fvcom_sshs["x2"].time.values[0]), str(fvcom_sshs["x2"].time.values[-1])
    )
    # NEMO sea surface height dataset
    try:
        nemo_ssh = _get_nemo_ssh(place, nemo_ssh_dataset_url_tmpl).sel(
            time=fvcom_period
        )
    except ValueError:
        # No NEMO sea surface height dataset for place
        nemo_ssh = None
    # CHS water level observations dataset
    try:
        obs_1min = (
            data_tools.get_chs_tides(
                "obs",
                place,
                arrow.get(fvcom_period.start) - timedelta(seconds=5 * 60),
                arrow.get(fvcom_period.stop),
            )
            .to_xarray()
            .rename({"index": "time"})
        )
        obs = obs_1min.to_dataset(name="water_level")
    except (AttributeError, KeyError):
        # No observations available, or invalid tide gauge station number
        obs_1min, obs = None, None
    # CHS water level predictions dataset
    try:
        pred_place = "Point Atkinson" if place == "Sandy Cove" else place
        try:
            pred = (
                data_tools.get_chs_tides(
                    "pred",
                    pred_place,
                    arrow.get(fvcom_period.start),
                    arrow.get(fvcom_period.stop),
                )
                .to_xarray()
                .rename({"index": "time"})
            )
            pred = pred.to_dataset(name="water_level")
        except (AttributeError, KeyError):
            # No predictions available, or invalid tide gauge station number
            pred = None
        # Residual differences between corrected model and observations and predicted tides
        fvcom_residuals = {}
        for model_config, fvcom_ssh in fvcom_sshs.items():
            try:
                fvcom_residuals[model_config] = fvcom_ssh - pred.water_level
                shared.localize_time(fvcom_residuals[model_config])
            except AttributeError:
                fvcom_residuals = None
        try:
            nemo_residual = nemo_ssh.ssh - pred.water_level
            shared.localize_time(nemo_residual)
        except AttributeError:
            nemo_residual = None
        try:
            obs_15min_avg = (obs_1min.resample(time="15min").mean())[1:]
            obs_15min = obs_15min_avg.to_dataset(name="water_level")
            obs_residual = obs_15min.water_level - pred.water_level
            shared.localize_time(obs_residual)
        except AttributeError:
            obs_residual = None
    except (TypeError, IndexError):
        # Invalid tide gauge station number, probably None
        pred, fvcom_residuals, nemo_residual, obs_residual = None, None, None, None
    # Change dataset times to Pacific time zone
    for fvcom_ssh in fvcom_sshs.values():
        shared.localize_time(fvcom_ssh)
    with suppress(IndexError, AttributeError):
        shared.localize_time(nemo_ssh)
    with suppress(IndexError, AttributeError):
        shared.localize_time(obs)
    with suppress(IndexError, AttributeError):
        shared.localize_time(pred)
    # Mean sea level
    msl = PLACES[place]["mean sea lvl"]
    return SimpleNamespace(
        fvcom_sshs=fvcom_sshs,
        nemo_ssh=nemo_ssh,
        obs=obs,
        pred=pred,
        msl=msl,
        fvcom_residuals=fvcom_residuals,
        nemo_residual=nemo_residual,
        obs_residual=obs_residual,
    )

def _get_nemo_ssh(place, dataset_url_tmpl):
    ## TODO: This is a work-around because neither netCDF4 nor xarray are able
    ##       to load the dataset directly from the URL due to an OpenDAP issue
    dataset_url = dataset_url_tmpl.format(place=place.replace(" ", ""))
    dataset_id = dataset_url.rsplit("/", 1)[1]
    ssh_file = Path("/tmp").joinpath(dataset_id).with_suffix(".nc")
    with ssh_file.open("wb") as f:
        resp = requests.get(f"{dataset_url}.nc")
        f.write(resp.content)
    try:
        nemo_ssh = xarray.open_dataset(ssh_file)
    except OSError:
        raise ValueError(f"NEMO ssh dataset not found for {place}")
    return nemo_ssh

def _prep_fig_axes(figsize, theme):
    fig, (ax_ssh, ax_res) = plt.subplots(
        nrows=2, figsize=figsize, facecolor=theme.COLOURS["figure"]["facecolor"]
    )
    ax_ssh.set_facecolor(theme.COLOURS["axes"]["background"])
    ax_res.set_facecolor(theme.COLOURS["axes"]["background"])
    register_matplotlib_converters()
    fig.autofmt_xdate()
    return fig, (ax_ssh, ax_res)

def _plot_water_level_time_series(ax, plot_data, theme):
    with suppress(AttributeError):
        # CHS sometimes returns an empty observations dataset
        if plot_data.obs.water_level.size:
            plot_data.obs.water_level.plot(
                ax=ax,
                linewidth=2,
                label="CHS Observed",
                color=theme.COLOURS["time series"]["tide gauge obs"],
            )
    with suppress(AttributeError):
        # CHS sometimes returns an empty prediction dataset
        if plot_data.pred.water_level.size:
            plot_data.pred.water_level.plot(
                ax=ax,
                linewidth=2,
                label="CHS Predicted",
                color=theme.COLOURS["time series"]["tidal prediction"],
                alpha=0.8,
            )
    with suppress(AttributeError):
        (plot_data.nemo_ssh.ssh + plot_data.msl).plot(
            ax=ax,
            linewidth=2,
            label="NEMO",
            color=theme.COLOURS["time series"]["tide gauge ssh"],
            alpha=0.8,
        )
    for model_config, fvcom_ssh in plot_data.fvcom_sshs.items():
        (fvcom_ssh + plot_data.msl).plot(
            ax=ax,
            linewidth=2,
            label=f"FVCOM {model_config.upper()}",
            color=theme.COLOURS["time series"]["vhfr fvcom ssh"][model_config],
            alpha=0.8,
        )

def _water_level_time_series_labels(ax, plot_data, place, theme):
    ax.set_title(
        f"Water Level at {place}",
        fontproperties=theme.FONTS["axes title"],
        color=theme.COLOURS["text"]["axes title"],
    )
    ax.set_xlabel("")
    ax.set_xlim(
        plot_data.fvcom_sshs["x2"].time.values[0],
        plot_data.fvcom_sshs["x2"].time.values[-1],
    )
    ax.set_ylabel(
        "Water Level above Chart Datum [m]",
        fontproperties=theme.FONTS["axis"],
        color=theme.COLOURS["text"]["axis"],
    )
    ax.legend(loc="best", prop=theme.FONTS["legend label small"])
    ax.grid(axis="both")
    theme.set_axis_colors(ax)

def _plot_residual_time_series(ax, plot_data, theme):
    with suppress(AttributeError):
        # CHS sometimes returns an empty observations dataset
        plot_data.obs_residual.plot(
            ax=ax,
            linewidth=2,
            label="CHS Observed",
            color=theme.COLOURS["time series"]["tide gauge obs"],
        )
    with suppress(TypeError):
        # No NEMO residual, probably because place is outside NEMO domaain
        (plot_data.nemo_residual + plot_data.msl).plot(
            ax=ax,
            linewidth=2,
            label="NEMO",
            color=theme.COLOURS["time series"]["tide gauge ssh"],
            alpha=0.8,
        )
    try:
        for model_config, fvcom_residual in plot_data.fvcom_residuals.items():
            (fvcom_residual + plot_data.msl).plot(
                ax=ax,
                linewidth=2,
                label=f"FVCOM {model_config.upper()}",
                color=theme.COLOURS["time series"]["vhfr fvcom ssh"][model_config],
                alpha=0.8,
            )
    except AttributeError:
        # No FVCOM residual, probably because no predicted water level
        # Plot invisible values so that we can label the x-axis, and add "Not Available" text
        ax.plot(
            plot_data.fvcom_sshs["x2"].time,
            numpy.zeros_like(plot_data.fvcom_sshs["x2"]),
            color=theme.COLOURS["axes"]["background"],
        )
        ax.text(
            0.5,
            0.5,
            "Not Available",
            fontproperties=theme.FONTS["axes annotation"],
            horizontalalignment="center",
            verticalalignment="center",
            transform=ax.transAxes,
        )

def _residual_time_series_labels(ax, plot_data, theme):
    ax.set_xlabel(
        f'Time [{plot_data.fvcom_sshs["x2"].attrs["tz_name"]}]',
        fontproperties=theme.FONTS["axis"],
        color=theme.COLOURS["text"]["axis"],
    )
    ax.set_xlim(
        plot_data.fvcom_sshs["x2"].time.values[0],
        plot_data.fvcom_sshs["x2"].time.values[-1],
    )
    ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter("%d%b %H:%M"))
    ax.set_ylabel(
        "Residual [m]",
        fontproperties=theme.FONTS["axis"],
        color=theme.COLOURS["text"]["axis"],
    )
    _, legend_labels = ax.get_legend_handles_labels()
    if legend_labels:
        # Don't try to draw a legend if there is no labels because doing so results in
        # WARNING:matplotlib.legend:No handles with labels found to put in legend.
        ax.legend(loc="best", prop=theme.FONTS["legend label small"])
    ax.grid(axis="both")
    theme.set_axis_colors(ax)


import logging, os, shutil, sys, re, scipy
from argparse import (ArgumentParser)
from mob_suite.version import __version__
import pandas as pd
import scipy.cluster.hierarchy as sch
from Bio import SeqIO

from scipy.cluster.hierarchy import fcluster
from mob_suite.blast import BlastRunner
from mob_suite.utils import \
    check_dependencies, \
    read_sequence_info, \
    read_file_to_dict, \
    read_fasta_dict, \
    NamesToTaxIDs

from mob_suite.wrappers import mash

from mob_suite.constants import LOG_FORMAT, ACS_LETTER_VALUES, ACS_FORMAT_VALUES, ACS_VALUES_TO_LETTERS, MAX_ACS_VALUE, \
    MOB_TYPER_REPORT_HEADER, MOB_CLUSTER_INFO_HEADER

def init_console_logger(lvl):
    logging_levels = [logging.ERROR, logging.WARN, logging.INFO, logging.DEBUG]
    report_lvl = logging_levels[lvl]

    logging.basicConfig(format=LOG_FORMAT, level=report_lvl)
    return logging

def parse_args():
    "Parse the input arguments, use '-h' for help"
    default_database_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'databases')
    parser = ArgumentParser(
        description="MOB-Cluster: Generate and update existing plasmid clusters' version: {}".format(__version__))
    parser.add_argument('-m', '--mode', type=str, required=True,
                        help='Build: Create a new database from scratch, Update: Update an existing database with one or more sequences')
    parser.add_argument('-f', '--infile', type=str, required=True,
                        help='Fasta file of sequences to cluster')
    parser.add_argument('-p', '--mob_typer_file', type=str, required=True,
                        help='MOB-typer report file for new sequences')
    parser.add_argument('-t', '--taxonomy', type=str, required=True,
                        help='TSV file for new sequences with the fields "id, organism"')
    parser.add_argument('-o', '--outdir', type=str, required=True, help='Output Directory to put results')
    parser.add_argument('-c', '--ref_cluster_file', type=str, required=False,
                        help='Existing MOB-cluster file to add the new sequences to')
    parser.add_argument('-r', '--ref_fasta_file', type=str, required=False,
                        help='Existing MOB-cluster fasta file of sequences contained in the MOB-cluster file')
    parser.add_argument('--num_threads', type=int, required=False, help='Number of threads to be used', default=1)
    parser.add_argument('--primary_cluster_dist', type=int, required=False,
                        help='Mash distance for assigning primary cluster id 0 - 1', default=0.06)
    parser.add_argument('--secondary_cluster_dist', type=int, required=False,
                        help='Mash distance for assigning primary cluster id 0 - 1', default=0.025)
    parser.add_argument('--debug', required=False, help='Show debug information', action='store_true')
    parser.add_argument('-V', '--version', action='version', version="%(prog)s " + __version__)
    return parser.parse_args()

'''
    Input: Accession string
    Return: True if string is in valid format : [A-Z][A-Z][0-9][0-9][0-9] else False

'''

def validate_acs_format(accession):
    if not isinstance(accession, str):
        logging.warning('Error provided accession number is not a string: {}'.format(accession))
        return False
    if not re.search("[A-Z][A-Z]\d\d\d", accession):
        logging.debug(
            'Error provided accession number is not in the correct format ->[A-Z][A-Z][0-9][0-9][0-9]<-: {}'.format(
                accession))
        return False
    else:
        return True

'''
    Input: Accession string ie. AA000
    Return: returns positive numerical integer value of Accession string, else returns -1

'''

def acs_to_int(accession):
    is_valid = validate_acs_format(accession)
    if not is_valid:
        logging.error('Error cannot continue due to invalid accession number {}'.format(accession))
        return -1

    return ACS_LETTER_VALUES[accession[0]] * ACS_FORMAT_VALUES[0] + \
           ACS_LETTER_VALUES[accession[1]] * ACS_FORMAT_VALUES[1] + \
           int(accession[2:5]) * ACS_FORMAT_VALUES[2]

'''
    Input: Integer
    Return: True if positive integer from 0 - MAX_ACS_VALUE

'''

def validate_int(number):
    if not isinstance(number, int):
        logging.error('Error provided numerical id is not a valid integer: {}'.format(number))
        return False
    if number < 0:
        logging.error('Error provided a negative number which is not a valid id: {}'.format(number))
        return False
    if number > MAX_ACS_VALUE:
        logging.error(
            'Error provided a number greater than what the existing ACS format can accommodate: {}'.format(number))
        return False
    else:
        return True

'''
    Input: integer of desired accession less than MAX Integer value the accession can accomodate
    Return: returns accession formated string or empty string on error

'''

def int_to_acs(numerical_id):
    is_valid = validate_int(numerical_id)
    if not is_valid:
        logging.error('Error cannot continue due to invalid id number {}'.format(numerical_id))
        return ""

    acs = [0, 0, 0]
    remainder = numerical_id
    for i in range(0, len(ACS_FORMAT_VALUES)):
        x = int(remainder / ACS_FORMAT_VALUES[i])
        remainder = remainder - (x * ACS_FORMAT_VALUES[i])
        acs[i] = x

    return "{}{}{}".format(ACS_VALUES_TO_LETTERS[acs[0]], ACS_VALUES_TO_LETTERS[acs[1]], str(acs[2]).zfill(3))

'''
    Input: Path to TSV file with MOB_CLUSTER_MINIMAL_INFO fields as the header lines
    Output: Dictionary of sequence indexed by sequence identifier
'''

def read_user_new_sequence_info(file):
    if os.path.getsize(file) == 0:
        return dict()
    data = pd.read_csv(file, sep='\t', header=0, names=MOB_TYPER_REPORT_HEADER, index_col=0)
    sequences = dict()
    header = list(data.head())

    for index, row in data.iterrows():
        if not index in sequences:
            sequences[index] = {}

        for i in range(0, len(header)):
            if row[header[i]] == 'nan':
                sequences[index][header[i]] = ''
            else:
                sequences[index][header[i]] = row[header[i]]

    return sequences

'''
    Input: fast file, mash reference sketch file, output file for mash distances
    Output: Mash distance results

'''

def calcDistances(input_fasta, ref_sketch_db, output_file):
    m = mash()
    mash_results = dict()
    with open(output_file, 'w', encoding="utf-8") as oh:
        m.run_mash(ref_sketch_db, input_fasta, oh)
        mash_lines = m.read_mash(output_file)
        for line in mash_lines:
            row = line.strip("\n").split("\t")
            ref_id = row[0]
            query_id = row[1]
            distance = float(row[2])
            if not query_id in mash_results:
                mash_results[query_id] = dict()
            mash_results[query_id][ref_id] = distance
    return mash_results

def write_clusters(file, cluster_assignments, header):
    with open(file, 'w', encoding="utf-8") as outfile:
        for id in cluster_assignments:
            outfile.write(str(id) + "\t" + "\t".join(cluster_assignments[id]))
        outfile.close()

def build_cluster_db(distance_matrix_file, distances):

    data = pd.read_csv(distance_matrix_file, sep='\t', header=0,
                       index_col=0)

    distance_matrix = data.to_numpy()
    condensed_matrix = scipy.spatial.distance.squareform(distance_matrix)
    Z = scipy.cluster.hierarchy.linkage(condensed_matrix, method='complete')

    clust_assignments = dict()

    for dist in distances:
        index = 0
        clusters = fcluster(Z, dist, criterion='distance')
        for id in data.columns.values:
            if not id in clust_assignments:
                clust_assignments[id] = list()
            clust_assignments[id].append(str(clusters[index]))
            index += 1
    return clust_assignments

def writeClusterAssignments(output_file, header, cluster_assignmnets):
    with open(output_file, 'w', encoding="utf-8") as out:
        out.write("\t".join(map(str, header)) + "\n")
        for id in cluster_assignmnets:
            line = []
            for h in header:
                if h == 'id':
                    continue
                if h in cluster_assignmnets[id]:
                    if cluster_assignmnets[id][h] == 'nan':
                        cluster_assignmnets[id][h] = '-'

                    line.append(str(cluster_assignmnets[id][h]))
                else:
                    line.append('')
            out.write("{}\n".format("\t".join(line)))
        out.close()

def appendFasta(new_seq_fasta, out_fasta):
    fh = open(out_fasta, 'a')
    for record in SeqIO.parse(new_seq_fasta, "fasta"):
        fh.write(">{}\n{}\n".format(record.id, record.seq))
    fh.close()

def updateFastaFile(in_fasta_file, out_fasta_file, cluster_assignments):
    out = open(out_fasta_file, 'w', encoding="utf-8")
    with open(in_fasta_file, "r") as handle:
        for record in SeqIO.parse(handle, "fasta"):
            row = str(record.id).split('|')
            id = row[0]

            if not id in cluster_assignments:
                continue

            out.write(">{}\n{}\n".format(id, record.seq))
        handle.close()
    out.close()

def selectCluster(clust_assignments, column):
    out = dict()
    for id in clust_assignments:
        row = str(id).split('|')
        out[row[0]] = clust_assignments[id][column]
    return out

def getMashDistances(mash_results_file, max_thresh=1):
    mashfile_handle = open(mash_results_file, 'r', encoding="utf-8")
    query_mash_distances = dict()

    for line in mashfile_handle:
        row = line.split('\t')
        query_id = row[1]
        ref_id = row[0].split('|')[0]
        distance = float(row[2])

        if distance > max_thresh:
            continue
        if not query_id in query_mash_distances:
            query_mash_distances[query_id] = dict()

        query_mash_distances[query_id][ref_id] = distance
    return query_mash_distances

def get_field_values(dictionary, field_key):
    values = []
    for key in dictionary:
        if field_key in dictionary[key]:
            values.append(dictionary[key][field_key])
    return values

def update_existing_db(new_seq_info, mob_cluster_seq_info, clust_assignments, primary_distance, secondary_distance,
                       num_threads=1):
    primary_acs = get_field_values(mob_cluster_seq_info, 'primary_cluster_id')
    primary_acs.sort(reverse=True)
    secondary_acs = get_field_values(mob_cluster_seq_info, 'secondary_cluster_id')
    secondary_acs.sort(reverse=True)

    acs = max(acs_to_int(primary_acs[0]), acs_to_int(secondary_acs[0])) + 1

    pri_acs_mapping = {}
    sec_acs_mapping = {}
    for seq_id in clust_assignments:
        p_acs = clust_assignments[seq_id][0]
        s_acs = clust_assignments[seq_id][1]
        if p_acs not in pri_acs_mapping:
            pri_acs_mapping[p_acs] = int_to_acs(acs)
            acs += 1
        if s_acs not in sec_acs_mapping:
            sec_acs_mapping[s_acs] = int_to_acs(acs)
            acs += 1

    for seq_id in new_seq_info:
        if seq_id in mob_cluster_seq_info:
            continue
        data = new_seq_info[seq_id]
        data['primary_dist'] = primary_distance
        data['secondary_dist'] = secondary_distance
        neighbor_id = data['mash_nearest_neighbor']
        dist = data['mash_neighbor_distance']
        if dist > primary_distance or neighbor_id not in mob_cluster_seq_info:
            data['primary_cluster_id'] = pri_acs_mapping[clust_assignments[seq_id][0]]
        else:
            data['primary_cluster_id'] = mob_cluster_seq_info[neighbor_id]['primary_cluster_id']

        if dist > secondary_distance or neighbor_id not in mob_cluster_seq_info:
            data['secondary_cluster_id'] = sec_acs_mapping[clust_assignments[seq_id][1]]
        else:
            data['secondary_cluster_id'] = mob_cluster_seq_info[neighbor_id]['secondary_cluster_id']

        mob_cluster_seq_info[seq_id] = data

    return mob_cluster_seq_info

'''
    Input: Two dictionaries
    Output: the key:value pairs which are not present in both dictionaries
'''

def find_different_dict_keys(first_dict, second_dict):
    set1 = first_dict.keys()
    set2 = second_dict.keys()

    return (set1 ^ set2)

'''
    Input: Dictionary indexed by md5 hash for sequences for new sequences and existing sequences.
    Output: Returns one sequence per md5 key which exists in the new set
'''

def getUniqSeqKeys(md5_new_seq_lookup, md5_mobcluster_seq_lookup):
    unique_keys = []

    for md5 in md5_new_seq_lookup:
        if md5 in md5_mobcluster_seq_lookup:
            continue
        for id in md5_new_seq_lookup[md5]:
            unique_keys.append(id)
            break

    return unique_keys

'''
    Input: Two dictionaries with identical sets of keys with seqStats produced by calcFastaStatsIndividual
    and userMeta produced by read_user_new_sequence_info
    Output: Combined dictionary with the elements from both dictionaries
'''

def combine_seqStats_with_userMeta(seqStats, userMeta):
    if len(seqStats.keys()) != len(userMeta):
        logging.error(
            'Error cannot continue due to difference number of sequences in fasta and metadata files: Num Fasta = {}, Num Meta = {}'.format(
                len(seqStats), len(userMeta)))
        return {}

    missing_keys = find_different_dict_keys(seqStats, userMeta)

    if len(missing_keys) > 0:
        logging.error(
            'Error cannot continue due to different ids in fasta and metadata files, keys are: '.format(missing_keys))
        return {}

    for id in userMeta:
        values = userMeta[id]
        for e in values:
            seqStats[id][e] = values[e]

    return seqStats

def convert_num_to_acs(clust_assignments):
    primary_keys = []
    secondary_keys = []

    for id in clust_assignments:

        p = clust_assignments[id][0]
        is_valid = validate_acs_format(p)
        if is_valid:
            p = int(acs_to_int(p))
        primary_keys.append(p)
        if len(clust_assignments[id]) > 1:
            s = clust_assignments[id][1]
            is_valid = validate_acs_format(s)
            if is_valid:
                s = int(acs_to_int(s))
            secondary_keys.append(s)
        clust_assignments[id][0] = int(p)
        clust_assignments[id][1] = int(s)

    primary_keys = list(set(primary_keys))
    unique_ints = []

    for p in primary_keys:
        is_valid = validate_acs_format(p)
        if is_valid:
            p = int(acs_to_int(p))
        unique_ints.append(int(p))

    highest_id = max(unique_ints)

    for id in clust_assignments:
        primary_acs = int_to_acs(clust_assignments[id][0])
        if len(clust_assignments[id]) > 1:
            secondary_acs = int_to_acs(clust_assignments[id][1] + highest_id)
        clust_assignments[id][0] = primary_acs
        clust_assignments[id][1] = secondary_acs

    return clust_assignments

def main():
    args = parse_args()
    if args.debug:
        logging = init_console_logger(3)
    else:
        logging = init_console_logger(2)
    logging.info('Running Mob-Suite Clustering toolkit v. {}'.format(__version__))
    logging.info('Processing fasta file {}'.format(args.infile))
    logging.info('Analysis directory {}'.format(args.outdir))

    check_dependencies(logging)

    input_fasta = args.infile
    if not os.path.isfile(input_fasta):
        logging.error('Error, input fasta specified does not exist: {}'.format(input_fasta))
        sys.exit()

    mob_typer_report_file = args.mob_typer_file
    if not os.path.isfile(mob_typer_report_file):
        logging.error('Error, input metadata file specified does not exist: {}'.format(mob_typer_report_file))
        sys.exit()

    mode = str(args.mode).lower()
    if mode not in ('update', 'build'):
        logging.error('Error you have not entered a valid mode of build or update, you entered: {}'.format(mode))
        sys.exit()

    out_dir = args.outdir
    num_threads = args.num_threads

    if not (args.primary_cluster_dist >= 0 and args.primary_cluster_dist <= 1):
        logging.error('Error distance thresholds must be between 0 - 1: {}'.format(args.primary_cluster_dist))
        sys.exit()
    else:
        primary_distance = args.primary_cluster_dist

    if not (args.secondary_cluster_dist >= 0 and args.secondary_cluster_dist <= 1):
        logging.error('Error distance thresholds must be between 0 - 1: {}'.format(args.secondary_cluster_dist))
        sys.exit()
    else:
        secondary_distance = args.secondary_cluster_dist

    if not os.path.isdir(out_dir):
        logging.info('Creating directory {}'.format(args.outdir))
        os.mkdir(out_dir, 0o755)
    tmp_dir = os.path.join(out_dir, '__tmp')
    if not os.path.isdir(tmp_dir):
        logging.info('Creating directory {}'.format(args.outdir))
        os.mkdir(tmp_dir, 0o755)

    taxonomy_file = args.taxonomy

    records = read_file_to_dict(mob_typer_report_file, MOB_TYPER_REPORT_HEADER, separater="\t")

    seq_ids = []
    new_seq_info = {}
    duplicate_keys = []
    for record in records:
        seq_ids.append(record['sample_id'])
        if not record['sample_id'] in new_seq_info:
            new_seq_info[record['sample_id']] = record
        else:
            duplicate_keys.append(record['sample_id'])

    if len(duplicate_keys) > 0:
        logging.error(
            "Duplicate sequence identifiers in fasta file. Please make every sequence id unique in the input file before using this tool")
        logging.error(
            "Duplicate sequence ids: {}".format(",".join(duplicate_keys)))
        sys.exit()

    record_identifications = read_file_to_dict(taxonomy_file, ['sample_id', 'organism'], separater="\t")
    organisms = []
    for record in record_identifications:
        organism = record['organism']
        if organism == 'unknown' or organism == '' or organism == 'Unknown':
            organism = 'Bacteria'
        organisms.append(organism)
        seq_id = record['sample_id']
        if seq_id in new_seq_info:
            new_seq_info[seq_id]['organism'] = organism

    taxids = NamesToTaxIDs(organisms)
    del(organisms)

    for seq_id in new_seq_info:
        organism = new_seq_info[seq_id]['organism']
        if organism in taxids:
            new_seq_info[seq_id]['taxid'] = taxids[organism][0]
        else:
            new_seq_info[seq_id]['taxid'] = 2

    if len(new_seq_info) == 0:
        logging.error('Error no MOB-typer results for sequences. Sequences must be typed with MOB-typer first')
        sys.exit()

    fasta_dict = read_fasta_dict(input_fasta)

    if len(fasta_dict) == 0:
        logging.error('Error no sequences found in input fasta: {}..cannot continue'.format(input_fasta))
        sys.exit()

    key_set_1 = set(seq_ids)
    key_set_2 = set(list(fasta_dict.keys()))

    if len(list(key_set_1 ^ key_set_2)) > 0:
        logging.error(
            'Error MOB-typer results: {} and input fasta: {} do not have the same set of identifiers, these must match in order to proceed'.format(
                mob_typer_report_file, input_fasta))
        logging.error(
            'Keys present in  MOB-typer results: {} and not in input fasta: {} are: {}'.format(mob_typer_report_file,
                                                                                input_fasta, list(
                    key_set_1 - key_set_2)))
        logging.error(
            'Keys present in  MOB-typer results: {} and not in input fasta: {} are: {}'.format(mob_typer_report_file,
                                                                                input_fasta, list(
                    key_set_2 - key_set_1)))
        sys.exit()

    tmp_cluster_file = os.path.join(out_dir, 'clusters.txt')
    tmp_ref_fasta_file = os.path.join(tmp_dir, 'references_tmp.fasta')
    update_fasta = os.path.join(out_dir, 'references_updated.fasta')

    # Sketch and calculate distances within update sequences
    if len(fasta_dict) > 1:
        mashObj = mash()
        mashObj.mashsketch(input_fasta, input_fasta + ".msh", num_threads=num_threads)
        distance_matrix_file = os.path.join(tmp_dir, 'mash_dist_matrix.txt')
        mashfile_handle = open(distance_matrix_file, 'w', encoding="utf-8")
        mashfile_handle.write(
            mashObj.run_mash(input_fasta + '.msh', input_fasta + '.msh', table=True, num_threads=num_threads).decode())
        mashfile_handle.close()
        clust_assignments = build_cluster_db(distance_matrix_file, (primary_distance, secondary_distance))
    else:
        seq_id = next(iter(fasta_dict))
        clust_assignments = {seq_id: [0, 1]}

    logging.info('Running MOB-cluster in {} mode'.format(mode))
    if mode == 'update':

        if args.ref_cluster_file is None:
            logging.error('Reference fasta file must be specified, please check help for parameter reference')
            sys.exit()

        ref_fasta = args.ref_fasta_file
        if not os.path.isfile(ref_fasta):
            logging.error('Reference fasta file specified does not exist: {}'.format(ref_fasta))
            sys.exit()

        if args.ref_cluster_file is None:
            logging.error('Reference cluster file must be specified, please check help for parameter reference')
            sys.exit()

        ref_cluster_file = args.ref_cluster_file
        if not os.path.isfile(ref_cluster_file):
            logging.error('Reference cluster file specified does not exist: {}'.format(ref_cluster_file))
            sys.exit()

        mob_cluster_seq_info = read_sequence_info(ref_cluster_file, MOB_CLUSTER_INFO_HEADER)

        logging.info('Running mob-cluster in update mode with input file: {}'.format(input_fasta))
        logging.info('Running mob-cluster in update mode with output directory: {}'.format(out_dir))
        logging.info('Running mob-cluster in update mode on reference fasta file: {}'.format(ref_fasta))
        logging.info('Reading previous cluster reference assignments from : {}'.format(ref_cluster_file))

        shutil.copy(ref_cluster_file, tmp_cluster_file)
        shutil.copy(ref_fasta, tmp_ref_fasta_file)
        logging.info('Creating new cluster assignments')
        new_seq_info = update_existing_db(new_seq_info,
                                          mob_cluster_seq_info,
                                          clust_assignments,
                                          primary_distance,
                                          secondary_distance,
                                          num_threads)

        cluster_assignments = {**mob_cluster_seq_info, **new_seq_info}
        logging.info('Writting cluster assignments to : {}'.format(tmp_cluster_file))
        writeClusterAssignments(tmp_cluster_file, MOB_CLUSTER_INFO_HEADER, cluster_assignments)
        shutil.copy(tmp_ref_fasta_file, os.path.join(out_dir, update_fasta))

    else:
        cluster_acs = convert_num_to_acs(clust_assignments)
        for id in cluster_acs:
            primary_key = cluster_acs[id][0]
            secondary_key = cluster_acs[id][1]
            new_seq_info[id]['primary_cluster_id'] = primary_key
            new_seq_info[id]['primary_dist'] = primary_distance
            new_seq_info[id]['secondary_cluster_id'] = secondary_key
            new_seq_info[id]['secondary_dist'] = secondary_distance

        writeClusterAssignments(tmp_cluster_file, MOB_CLUSTER_INFO_HEADER, new_seq_info)
        shutil.copy(input_fasta, update_fasta)

    logging.info("Sketching new fasta {}".format(update_fasta))
    mash_db_file = "{}.msh".format(update_fasta)
    mObj = mash()
    mObj.mashsketch(update_fasta, mash_db_file, num_threads=num_threads)
    logging.info("Building blastdb {}".format(update_fasta))
    blast_runner = BlastRunner(update_fasta, '')
    blast_runner.makeblastdb(update_fasta, 'nucl', logging=logging)
    logging.info("Removing temporary directory")
    shutil.rmtree(tmp_dir)
    logging.info("MOB-cluster completed, analysis results written to {}".format(out_dir))

# call main function
if __name__ == '__main__':
    main()