import numpy as np
import torch
from gensim.models import word2vec
from transformers import BertTokenizer, BertModel

from dataset import get_vocab

class Word2VecEmbedding:
    def __init__(self, word2vec_file: str):
        self.embedding_type = 'word2vec'
        model = word2vec.Word2Vec.load(word2vec_file)
        embedding_dim = model.vector_size
        model.wv["<pad>"] = np.full((1, embedding_dim), 0)
        model.wv["<unk>"] = np.random.rand(embedding_dim)
        self.model = model.wv
        self.key_to_idx = model.wv.key_to_index
        self.embedding_dim = embedding_dim
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    def get_embeddings(self, text):
        embedded = [[self.model[word] for word in sentence.split(' ')] for sentence in text]
        embedded = torch.FloatTensor(np.array(embedded)).to(self.device)
        embedded = torch.permute(embedded, [1, 0, 2])
        return embedded

class BERTEmbedding:

    def __init__(self, vocab_file: str):
        self.embedding_type = 'bert'
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        word2idx, idx2word, vocab_size = get_vocab(vocab_file)
        all_vocab = list(word2idx.keys())
        special_tokens = all_vocab[:2]
        vocab_tokens = all_vocab[2:]
        self.tokenizer.add_tokens(special_tokens, special_tokens=True)

        self.tokenizer.add_tokens(vocab_tokens)

        self.embedding_dim = 768

        self.bert = BertModel.from_pretrained('bert-base-uncased', output_hidden_states=True).to(self.device)
        self.bert.resize_token_embeddings(len(self.tokenizer))
        self.bert.eval()

    def get_embeddings(self, text):
        tokenized_texts = [self.tokenizer.tokenize(sentence) for sentence in text]

        indexed_tokens = [self.tokenizer.convert_tokens_to_ids(tokenized_text) for tokenized_text in tokenized_texts]
        segments_ids = [[1] * len(tokenized_text) for tokenized_text in tokenized_texts]

        token_tensor = torch.tensor(indexed_tokens).to(self.device)
        segment_tensor = torch.tensor(segments_ids).to(self.device)

        with torch.no_grad():
            outputs = self.bert(token_tensor, segment_tensor)
            hidden_states = outputs[2]
            token_embeddings = torch.stack(hidden_states, dim=0)
            second_to_last_layer = token_embeddings[11]
            embedding_tensor = torch.permute(second_to_last_layer, [1, 0, 2])
            return embedding_tensor


import gym
import sys
import os
import itertools
import collections
import numpy as np
import tensorflow as tf

from inspect import getsourcefile
current_path = os.path.dirname(os.path.abspath(getsourcefile(lambda: 0)))
import_path = os.path.abspath(os.path.join(current_path, "../.."))

if import_path not in sys.path:
    sys.path.append(import_path)

# from lib import plotting
from lib.atari.state_processor import StateProcessor
from lib.atari import helpers as atari_helpers
from estimators import ValueEstimator, PolicyEstimator

Transition = collections.namedtuple(
    "Transition", ["state", "action", "reward", "next_state", "done"])

def make_copy_params_op(v1_list, v2_list):
    """
    Creates an operation that copies parameters from variable in v1_list to variables in v2_list.
    The ordering of the variables in the lists must be identical.
    """
    v1_list = list(sorted(v1_list, key=lambda v: v.name))
    v2_list = list(sorted(v2_list, key=lambda v: v.name))

    update_ops = []
    for v1, v2 in zip(v1_list, v2_list):
        op = v2.assign(v1)
        update_ops.append(op)

    return update_ops

def make_train_op(local_estimator, global_estimator):
    """
    Creates an op that applies local estimator gradients
    to the global estimator.
    """
    local_grads, _ = zip(*local_estimator.grads_and_vars)
    # Clip gradients
    local_grads, _ = tf.clip_by_global_norm(local_grads, 5.0)
    _, global_vars = zip(*global_estimator.grads_and_vars)
    local_global_grads_and_vars = list(zip(local_grads, global_vars))
    return global_estimator.optimizer.apply_gradients(local_global_grads_and_vars,
                                                      global_step=tf.contrib.framework.get_global_step())

class Worker(object):
    """
    An A3C worker thread. Runs episodes locally and updates global shared value and policy nets.

    Args:
      name: A unique name for this worker
      env: The Gym environment used by this worker
      policy_net: Instance of the globally shared policy net
      value_net: Instance of the globally shared value net
      global_counter: Iterator that holds the global step
      discount_factor: Reward discount factor
      summary_writer: A tf.train.SummaryWriter for Tensorboard summaries
      max_global_steps: If set, stop coordinator when global_counter > max_global_steps
    """

    def __init__(self, name, env, policy_net, value_net, global_counter, discount_factor=0.99, summary_writer=None, max_global_steps=None):
        self.name = name
        self.discount_factor = discount_factor
        self.max_global_steps = max_global_steps
        self.global_step = tf.contrib.framework.get_global_step()
        self.global_policy_net = policy_net
        self.global_value_net = value_net
        self.global_counter = global_counter
        self.local_counter = itertools.count()
        self.sp = StateProcessor()
        self.summary_writer = summary_writer
        self.env = env

        # Create local policy/value nets that are not updated asynchronously
        with tf.variable_scope(name):
            self.policy_net = PolicyEstimator(policy_net.num_outputs)
            self.value_net = ValueEstimator(reuse=True)

        # Op to copy params from global policy/valuenets
        self.copy_params_op = make_copy_params_op(
            tf.contrib.slim.get_variables(
                scope="global", collection=tf.GraphKeys.TRAINABLE_VARIABLES),
            tf.contrib.slim.get_variables(scope=self.name, collection=tf.GraphKeys.TRAINABLE_VARIABLES))

        self.vnet_train_op = make_train_op(
            self.value_net, self.global_value_net)
        self.pnet_train_op = make_train_op(
            self.policy_net, self.global_policy_net)

        self.state = None

    def run(self, sess, coord, t_max):
        with sess.as_default(), sess.graph.as_default():
            # Initial state
            self.state = atari_helpers.atari_make_initial_state(
                self.sp.process(self.env.reset()))
            try:
                while not coord.should_stop():
                    # Copy Parameters from the global networks
                    sess.run(self.copy_params_op)

                    # Collect some experience
                    transitions, local_t, global_t = self.run_n_steps(
                        t_max, sess)

                    if self.max_global_steps is not None and global_t >= self.max_global_steps:
                        tf.logging.info(
                            "Reached global step {}. Stopping.".format(global_t))
                        coord.request_stop()
                        return

                    # Update the global networks
                    self.update(transitions, sess)

            except tf.errors.CancelledError:
                return

    def _policy_net_predict(self, state, sess):
        feed_dict = {self.policy_net.states: [state]}
        preds = sess.run(self.policy_net.predictions, feed_dict)
        return preds["probs"][0]

    def _value_net_predict(self, state, sess):
        feed_dict = {self.value_net.states: [state]}
        preds = sess.run(self.value_net.predictions, feed_dict)
        return preds["logits"][0]

    def run_n_steps(self, n, sess):
        transitions = []
        for _ in range(n):
            # Take a step
            if np.random.uniform() < 0.1:
                action = np.random.randint(0, self.env.actions_state.n)
            else:
                action_probs = self._policy_net_predict(self.state, sess)
                action = np.random.choice(np.arange(len(action_probs)), p=action_probs)
            next_state, reward, done, _ = self.env.step(action)
            next_state = atari_helpers.atari_make_next_state(
                self.state, self.sp.process(next_state))

            # Store transition
            transitions.append(Transition(
                state=self.state, action=action, reward=reward, next_state=next_state, done=done))

            # Increase local and global counters
            local_t = next(self.local_counter)
            global_t = next(self.global_counter)

            if local_t % 100 == 0:
                tf.logging.info("{}: local Step {}, global step {}".format(
                    self.name, local_t, global_t))

            if done:
                self.state = atari_helpers.atari_make_initial_state(
                    self.sp.process(self.env.reset()))
                break
            else:
                self.state = next_state
        return transitions, local_t, global_t

    def update(self, transitions, sess):
        """
        Updates global policy and value networks based on collected experience

        Args:
          transitions: A list of experience transitions
          sess: A Tensorflow session
        """

        # If we episode was not done we bootstrap the value from the last state
        reward = 0.0
        if not transitions[-1].done:
            reward = self._value_net_predict(transitions[-1].next_state, sess)

        # Accumulate minibatch exmaples
        states = []
        policy_targets = []
        value_targets = []
        actions = []

        for transition in transitions[::-1]:
            reward = transition.reward + self.discount_factor * reward
            policy_target = (
                reward - self._value_net_predict(transition.state, sess))
            # Accumulate updates
            states.append(transition.state)
            actions.append(transition.action)
            policy_targets.append(policy_target)
            value_targets.append(reward)

        feed_dict = {
            self.policy_net.states: np.array(states),
            self.policy_net.targets: policy_targets,
            self.policy_net.actions: actions,
            self.value_net.states: np.array(states),
            self.value_net.targets: value_targets,
        }

        # Train the global estimators using local gradients
        global_step, pnet_loss, vnet_loss, _, _, pnet_summaries, vnet_summaries = sess.run([
            self.global_step,
            self.policy_net.loss,
            self.value_net.loss,
            self.pnet_train_op,
            self.vnet_train_op,
            self.policy_net.summaries,
            self.value_net.summaries
        ], feed_dict)

        # Write summaries
        if self.summary_writer is not None:
            self.summary_writer.add_summary(pnet_summaries, global_step)
            self.summary_writer.add_summary(vnet_summaries, global_step)
            self.summary_writer.flush()

        return pnet_loss, vnet_loss, pnet_summaries, vnet_summaries


from collections import OrderedDict

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim

from rl.dataset import ReplayBuffer, MultiProcessReplayBuffer
from rl.agents.base_agent import BaseAgent
from util.logger import logger
from util.pytorch import optimizer_cuda, count_parameters, \
    obs2tensor, to_tensor, list2dict

class PPOAgent(BaseAgent):
    def __init__(self, config, ob_space, ac_space,
                 actor, critic):
        super().__init__(config, ob_space)

        self._ac_space = ac_space
        # build up networks
        self._actor = actor(config, ob_space, ac_space, config.tanh_policy, activation='tanh')
        self._old_actor = actor(config, ob_space, ac_space, config.tanh_policy, activation='tanh')
        self._critic = critic(config, ob_space, activation='tanh')
        self._network_cuda(config.device)

        self._actor_optim = optim.Adam(self._actor.parameters(), lr=config.lr_actor)
        self._critic_optim = optim.Adam(self._critic.parameters(), lr=config.lr_critic)

        # self._buffer = ReplayBuffer(['ob', 'ac', 'done', 'rew', 'ret', 'adv', 'ac_before_activation'],
        self._buffer = MultiProcessReplayBuffer(config,
                                    config.num_processes,
                                    ob_space, ac_space)

        logger.info('Creating a PPO agent')
        logger.info('The actor has %d parameters', count_parameters(self._actor))
        logger.info('The critic has %d parameters', count_parameters(self._critic))

    def store_episode(self, rollouts):
        self._compute_gae(rollouts)
        self._buffer.store_episode(rollouts)

    def store_sample(self, rollout):
        self._buffer.store_sample(rollout)

    def _compute_gae(self, rollouts):
        """ Computes GAE from @rollouts. """
        config = self._config
        T = len(rollouts['done'])
        vpred = np.array(rollouts['vpred']).squeeze(-1)
        assert len(vpred) == T + 1

        done = rollouts['done']
        rew = rollouts['rew']
        adv = np.empty((T, config.num_processes) , 'float32')
        lastgaelam = 0
        for t in reversed(range(T)):
            nonterminal = 1 - done[t]
            delta = rew[t] + self._config.discount_factor * vpred[t + 1] * nonterminal - vpred[t]
            adv[t] = lastgaelam = delta + self._config.discount_factor * self._config.gae_lambda * nonterminal * lastgaelam

        ret = adv + vpred[:-1]

        assert np.isfinite(adv).all()
        assert np.isfinite(ret).all()

        # update rollouts
        if adv.std() == 0:
            rollouts['adv'] = (adv * 0)
        else:
            rollouts['adv'] = ((adv - adv.mean()) / adv.std())
        rollouts['ret'] = ret

    def state_dict(self):
        return {
            'actor_state_dict': self._actor.state_dict(),
            'critic_state_dict': self._critic.state_dict(),
            'actor_optim_state_dict': self._actor_optim.state_dict(),
            'critic_optim_state_dict': self._critic_optim.state_dict(),
            'ob_norm_state_dict': self._ob_norm.state_dict(),
        }

    def load_state_dict(self, ckpt):
        self._actor.load_state_dict(ckpt['actor_state_dict'])
        self._critic.load_state_dict(ckpt['critic_state_dict'])
        self._ob_norm.load_state_dict(ckpt['ob_norm_state_dict'])
        self._network_cuda(self._config.device)

        self._actor_optim.load_state_dict(ckpt['actor_optim_state_dict'])
        self._critic_optim.load_state_dict(ckpt['critic_optim_state_dict'])
        optimizer_cuda(self._actor_optim, self._config.device)
        optimizer_cuda(self._critic_optim, self._config.device)

    def _network_cuda(self, device):
        self._actor.to(device)
        self._old_actor.to(device)
        self._critic.to(device)

    def train(self):
        self._soft_update_target_network(self._old_actor, self._actor, 0.0)
        for _ in range(self._config.ppo_epoch):
            generator = self._buffer.generator(self._config.batch_size)
            for transitions in generator:
                train_info = self._update_network(transitions)

        self._buffer.clear()

        return train_info

    def _update_network(self, transitions):
        info = {}

        # pre-process observations
        o = transitions['ob']
        # o = self.normalize(o)
        bs = len(transitions['done'])
        _to_tensor = lambda x: to_tensor(x, self._config.device)
        o = _to_tensor(o)
        ac = _to_tensor(transitions['ac'])
        a_z = _to_tensor(transitions['ac_before_activation'])
        ret = _to_tensor(transitions['ret']).reshape(bs, 1)
        adv = _to_tensor(transitions['adv']).reshape(bs, 1)
        vpred = _to_tensor(transitions['vpred']).reshape(bs, 1)
        old_log_pi = _to_tensor(transitions['log_prob']).reshape(bs, 1)

        log_pi, ent = self._actor.act_log(o, a_z)
        # old_log_pi, _ = self._old_actor.act_log(o, a_z)
        # if old_log_pi.min() < -100:
        #     import ipdb; ipdb.set_trace()

        # the actor loss
        entropy_loss = self._config.entropy_loss_coeff * ent.mean()
        ratio = torch.exp(log_pi - old_log_pi.detach())
        surr1 = ratio * adv
        surr2 = torch.clamp(ratio, 1.0 - self._config.clip_param,
                            1.0 + self._config.clip_param) * adv
        actor_loss = -torch.min(surr1, surr2).mean()

        if not np.isfinite(ratio.cpu().detach()).all() or not np.isfinite(adv.cpu().detach()).all():
            import ipdb; ipdb.set_trace()
        info['entropy_loss'] = entropy_loss.cpu().item()
        info['actor_loss'] = actor_loss.cpu().item()
        actor_loss += entropy_loss

        # the q loss
        value_pred = self._critic(o)
        value_pred_clipped = vpred + (value_pred - vpred).clamp(-self._config.clip_param, self._config.clip_param)
        value_loss = (value_pred-ret).pow(2)
        value_loss_clipped = (value_pred_clipped-ret).pow(2)
        value_loss = self._config.value_loss_coeff * torch.max(value_loss, value_loss_clipped).mean()

        # value_loss = self._config.value_loss_coeff * (ret - value_pred).pow(2).mean()

        info['value_target'] = ret.mean().cpu().item()
        info['value_predicted'] = value_pred.mean().cpu().item()
        info['value_loss'] = value_loss.cpu().item()

        # update the actor
        self._actor_optim.zero_grad()
        actor_loss.backward()
        if self._config.max_grad_norm is not None:
            torch.nn.utils.clip_grad_norm_(self._actor.parameters(), self._config.max_grad_norm)
        self._actor_optim.step()

        # update the critic
        self._critic_optim.zero_grad()
        value_loss.backward()
        if self._config.max_grad_norm is not None:
            torch.nn.utils.clip_grad_norm_(self._critic.parameters(), self._config.max_grad_norm)
        self._critic_optim.step()

        # include info from policy
        info.update(self._actor.info)

        return info


import numpy as np
import pandas as pd
from mm_utils import oneHot, read_genome
from models import pearson
from mm_gen import GenScan
from tensorflow.keras.models import load_model
import argparse
import re
import tensorflow.keras.backend as K
import tensorflow as tf

def user_input():
    parser = argparse.ArgumentParser()

    parser.add_argument('-f', '--flank_size', help='Number of nt\'s to pad the sequence on each side',
                        type=int, default=100, required=False)
    parser.add_argument('-bs', '--batch_size', help='Batch size', type=int, required=False)
    parser.add_argument('-dp', '--data_path', help='Path to chromosome', type=str, required=False)
    parser.add_argument('-of', '--out_file', help='Path for the outputs', type=str,
                        required=False, default='.')
    parser.add_argument('-mp', '--model_path', help='Path to pre-trained model.', type=str, required=False)
    parser.add_argument('-rc', '--reverse_complement', help='Use reverse complement of chromosome.', type=int,
                        required=False, default=0, choices=[0, 1])
    parser.add_argument('-c', '--cut', help='Cut N tracts at the ends (for chromosome scanning).', type=int,
                        required=False, default=0, choices=[0, 1])
    parser.add_argument('-obg', '--out_bg', help='Name of output bedGraph file', type=str, required=False,
                        default='scan_res')
    parser.add_argument('-a', '--append',
                        help='If 1 append every sequence with flanks of Ns of length \"flank\", if 0 leave as is. '
                             '\nIf not used, the first and last 100 nt are considered flanks.',
                        type=int, required=False, default=0)
    parser.add_argument('-w', '--workers', help='Number of workers for batch generation.', type=int,
                        required=False, default=10)
    parser.add_argument('-fb', '--filter_bins',
                        help='If passed, bins with predicted mm score under the input value would be filtered out.',
                        type=float, required=False)
    parser.add_argument('-m', '--merge', help='Merge consecutive bins.',
                        type=int, required=False, default=0, choices=[0, 1])
    parser.add_argument('-q', '--max_queue_size', help='Max queu size.', type=int,
                        required=False, default=20)
    parser.add_argument('-mpr', '--multiprocessing', help='Use multiprocessing.', type=int,
                        required=False, default=1, choices=[0, 1])
    parser.add_argument('-vn', '--val_n',
                        help='Encoding of base N.',
                        type=float, required=False, default=0.25)

    args = parser.parse_args()
    arguments = vars(args)

    return arguments

def main(args):
    gpus = tf.config.experimental.list_physical_devices('GPU')
    if gpus:
        try:
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
        except RuntimeError as e:
            print(e)

    # tf.config.set_visible_devices([], 'GPU')

    win_size = 15
    flank = args['flank_size']

    model = load_model(args['model_path'], custom_objects={"pearson": pearson}, compile=False)

    seqs = read_genome(args['data_path'])

    df = pd.DataFrame(columns=['name', 'start', 'end', 'pred'])

    s = 0

    for seq in seqs:
        # remove large N strands from both ends, but leave flanked windows on each side
        c = str(seq.seq)
        if args['append']:
            c = 'N' * flank + c + 'N' * flank
        if args['cut']:
            try:
                s = re.search(r'NNNNNN[ACGT]', c).end() - 1 - flank
            except:
                s = 0
            try:
                e = len(c) - re.search(r'NNNNNN[ACGT]', c[::-1]).end() + flank + 1
            except:
                e = len(c)
            print('cutting chromosome {} ends at {:d} and {:d}'.format(seq.name, s, e))
            c = c[s:e]

        c = oneHot(c, val_N=args['val_n'])

        if args['batch_size']:
            bs = args['batch_size']
        else:
            bs = len(c) - flank * 2 - win_size

        scan_gen = GenScan(win_size, bs=bs, seq=c, flank=flank, rc=args['reverse_complement'])

        pred = model.predict(scan_gen, use_multiprocessing=bool(args['multiprocessing']), workers=args['workers'],
                             max_queue_size=args['max_queue_size'])
        pred = pred.squeeze()
        num_win = len(pred)

        if args['merge']:
            if args['filter_bins'] is not None:
                tmp_df = pd.DataFrame(
                    {'name': [seq.name] * num_win,
                     'start': s + flank + np.arange(num_win),
                     'end': s + flank + np.arange(num_win) + win_size,
                     'pred': pred})
                tmp_df = tmp_df[tmp_df['pred'] >= args['filter_bins']]
                tmp_df.reset_index(inplace=True, drop=True)
                split_co = tmp_df['end'][:-1].reset_index(drop=True) - tmp_df['start'][1:].reset_index(drop=True)
                split_co[split_co > -1] = 0
                split_co = pd.concat([pd.Series([0]), split_co]).reset_index(drop=True)
                split = np.split(tmp_df, np.flatnonzero(split_co))
                tmp_df = pd.DataFrame(columns=['name', 'start', 'end', 'pred'])
                col_names = list(tmp_df)
                data = []
                for x in split:
                    l = [x['name'].iloc[0], x['start'].iloc[0], x['end'].iloc[-1], x['pred'].max()]
                    d = dict(zip(col_names, l))
                    data.append(d)
                tmp_df = tmp_df.append(data, True)
            else:
                tmp_df = pd.DataFrame(
                    {'name': [seq.name],
                     'start': [s + flank],
                     'end': [s + flank + num_win + win_size],
                     'pred': [pred.max()]})

        else:
            tmp_df = pd.DataFrame(
                {'name': [seq.name] * num_win,
                 'start': s + np.arange(num_win),
                 'end': s + np.arange(num_win) + win_size,
                 'pred': [pred]})

        df = pd.concat([df, tmp_df])

        K.clear_session()

    df.to_csv("{}/{}.bedGraph".format(args['out_file'], args['out_bg']),
              index=False, header=False, sep='\t', float_format='%.1f')

if __name__ == "__main__":
    args = user_input()
    main(args)


import sys
from os.path import abspath, dirname, join
sys.path.insert(1, abspath(join(dirname(dirname(__file__)), 'src')))

import unittest
import numpy as np
import cv2 as cv

import face

class TestLiviness(unittest.TestCase):
    def setUp(self) -> None:
        self.num_landmarks_points = 468
        self.dims = (640, 480)

        self.frame = np.random.random((self.dims[0], self.dims[1], 3)) * 255
        self.frame = self.frame.astype('uint8')

    def test_landmarks_detection(self):
        results = face.liveness.landmarks_detection(frame=self.frame, draw=False)
        self.assertIsInstance(results, tuple)
        self.assertIsInstance(results[0], list)
        self.assertIsInstance(results[1], bool)

    def test_blink(self):
        landmarks = [(0, 0)] * self.num_landmarks_points
        blink_ratio = face.liveness.blink_detection(landmarks, face.config.RIGHT_EYE, face.config.LEFT_EYE)
        self.assertIsInstance(blink_ratio, float)
        self.assertEqual(blink_ratio, 0.0)

    def test_blur(self):
        blur_ratio = face.liveness.blur_detection(frame=self.frame)
        self.assertIsInstance(blur_ratio, np.float64)

    def test_lbp(self):
        gray_frame = cv.cvtColor(self.frame, cv.COLOR_BGR2GRAY)
        lbp = face.liveness.LocalBinaryPatterns(24, 8)
        lbp_ratio = lbp.get_lbp_max(gray_frame)
        self.assertIsInstance(lbp_ratio, np.float64)


import logging
import operator as op
from collections.abc import Sequence
from copy import copy, deepcopy
from datetime import datetime
from os.path import splitext
import inspect

from tqdm import tqdm
import astropy.constants as const
import astropy.units as u
import numpy as np
import specutils
import specutils.manipulation as specman
from astropy import coordinates as coords
from astropy.io import fits
from astropy.time import Time
from astropy.nddata import StdDevUncertainty

import astroplan

from . import reference_frame as rf
from .data_modules.stellar_db import StellarDb
from .simulator import detector
from .extractor.steps import StepIO

from flex.flex import FlexFile
from flex.extensions.bindata import MultipleDataExtension

logger = logging.getLogger(__name__)

class SpectrumBase:
    reference_frame_values = ["barycentric", "telescope", "planet", "star"]

    @property
    def datetime(self):
        return self.meta["datetime"]

    @datetime.setter
    def datetime(self, value):
        if not isinstance(value, Time):
            value = Time(value)
        self.meta["datetime"] = value

    @property
    def reference_frame(self):
        return self.meta["reference_frame"]

    @reference_frame.setter
    def reference_frame(self, value):
        if (
            not isinstance(value, rf.ReferenceFrame)
            and value not in self.reference_frame_values
        ):
            raise ValueError(
                f"Reference frame not understood."
                f"Expected one of {self.reference_frame_values} but got {value}"
            )
        if value in self.reference_frame_values:
            value = self.reference_frame_from_name(value)
        self.meta["reference_frame"] = value

    def reference_frame_from_name(self, frame):
        frame = rf.reference_frame_from_name(
            frame,
            star=self.meta["star"],
            planet=self.meta["planet"],
            observatory=self.meta["observatory_location"],
        )
        return frame

    @property
    def star(self):
        return self.meta["star"]

    @star.setter
    def star(self, value):
        self.meta["star"] = value

    @property
    def planet(self):
        return self.meta["planet"]

    @planet.setter
    def planet(self, value):
        self.meta["planet"] = value

    @property
    def observatory_location(self):
        return self.meta["observatory_location"]

    @observatory_location.setter
    def observatory_location(self, value):
        self.meta["observatory_location"] = value

    @property
    def airmass(self):
        target = astroplan.FixedTarget(name=self.star.name, coord=self.star.coordinates)
        observer = astroplan.Observer(self.observatory_location)
        altaz = observer.altaz(self.datetime, target)
        airmass = altaz.secz.value
        return airmass

    def __flex_save__(self):
        header = self.meta
        header["wavelength_unit"] = self.wavelength.unit
        header["flux_unit"] = self.flux.unit
        data = {
            "wavelength": self.wavelength.to_value("AA"),
            "flux": self.flux.to_value(1),
        }
        if self.uncertainty is not None:
            header["uncertainty_unit"] = self.uncertainty.unit
            data["uncertainty"] = self.uncertainty.array

        module = self.__class__.__module__
        cls = self.__class__.__name__
        ext = MultipleDataExtension(header, data, cls=f"{module}.{cls}")
        return ext

    @classmethod
    def __flex_load__(cls, header, data):
        ext = MultipleDataExtension._parse(header, data)

        wavelength = ext.data["wavelength"]
        wavelength = wavelength << ext.header["wavelength_unit"]

        flux = ext.data["flux"]
        flux = flux << ext.header["flux_unit"]

        if "uncertainty" in ext.data:
            uncertainty = ext.data["uncertainty"]
            uncertainty = uncertainty << ext.header["uncertainty_unit"]
            uncertainty = StdDevUncertainty(uncertainty)
        else:
            uncertainty = None

        exceptions = [
            "__module__",
            "__class__",
            "__header__",
            "uncertainty_unit",
            "wavelength_unit",
            "flux_unit",
        ]
        meta = {k: v for k, v in ext.header.items() if k not in exceptions}
        meta["spectral_axis"] = wavelength
        meta["flux"] = flux
        meta["uncertainty"] = uncertainty

        return cls(**meta)

    def to_dict(self):
        data = copy(self.meta)
        data["spectral_axis"] = self.wavelength.value
        data["spectral_axis_unit"] = self.wavelength.unit
        data["flux"] = self.flux.value
        data["flux_unit"] = self.flux.unit
        if self.uncertainty is not None:
            data["uncertainty"] = self.uncertainty.array
            data["uncertainty_unit"] = self.uncertainty.unit
        return data

    @classmethod
    def from_dict(cls, data):
        meta = data.get("meta", {})
        exceptions = ["flux", "flux_unit", "spectral_axis", "spectral_axis_unit"]
        for k, v in data.items():
            if k not in exceptions:
                meta[k] = v
        flux_unit = data["flux_unit"]
        wave_unit = data["spectral_axis_unit"]
        meta["flux"] = data["flux"] << u.Unit(flux_unit)
        meta["spectral_axis"] = data["spectral_axis"] << u.Unit(wave_unit)

        if "uncertainty" in data.keys():
            uncs_unit = data["uncertainty_unit"]
            uncs = data["uncertainty"] << u.Unit(uncs_unit)
            uncs = StdDevUncertainty(uncs)
            meta["uncertainty"] = uncs

        self = cls(**meta)
        return self

    def get_fits_hdu(self):
        raise NotImplementedError

    @classmethod
    def read_fits_hdu(cls, hdu):
        raise NotImplementedError

    @staticmethod
    def determine_filetype_based_on_filename(filename):
        _, ext = splitext(filename)
        if ext in [".flex", ".flx"]:
            return "flex"
        elif ext in [".fits", ".gz"]:
            return "fits"
        elif ext in [".npz"]:
            return "npz"
        else:
            raise ValueError("Could not determine filetype based on file ending")

    def write(self, filename, format=None):
        if format is None:
            format = self.__class__.determine_filetype_based_on_filename(filename)
        if format == "flex":
            ff = FlexFile(extensions={"spectrum": self.__flex_save__()})
            ff.write(filename)
        elif format == "fits":
            hdu = self.get_fits_hdu()
            hdu.writeto(filename, overwrite=True)
        elif format == "npz":
            data = self.to_dict()
            np.savez(filename, **data)
        else:
            raise ValueError(
                f"Format not understood expected one of ('flex', 'fits') but got {format}"
            )

    @classmethod
    def read(cls, filename, format=None):
        if format is None:
            format = cls.determine_filetype_based_on_filename(filename)

        if format == "flex":
            ff = FlexFile.read(filename)
            spec = ff["spectrum"]
            return spec
        elif format == "fits":
            hdulist = fits.open(filename)
            spec = cls.read_fits_hdu(hdulist[1])
            return spec
        elif format == "npz":
            data = np.load(filename, allow_pickle=True)
            data = {k: v[()] for k, v in data.items()}
            spec = cls.from_dict(data)
            return spec
        else:
            raise ValueError(
                f"Format not understood expected one of ('flex', 'fits') but got {format}"
            )

class Spectrum1D(SpectrumBase, specutils.Spectrum1D):
    """
    Extends the specutils Spectrum1D class
    with a few convenience functions useful for
    our case.
    Most importently for resampling and shifting
    the spectrum to different wavelengths grids
    """

    __code__ = [
        m
        for m in inspect.getmembers(specutils.Spectrum1D.__init__)
        if m[0] == "__code__"
    ][0][1]
    __init_args__ = inspect.getargs(__code__).args[1:] + [
        "data",
        "unit",
        "uncertainty",
        "meta",
        "mask",
        "copy",
    ]

    def __init__(self, *args, **kwargs):
        # Set default options (if not given)

        if "spectral_axis" not in kwargs.keys():
            kwargs["radial_velocity"] = kwargs.get("radial_velocity", 0 * (u.km / u.s))

        meta = {}
        # Which data source was this obtained from
        meta["source"] = kwargs.pop("source", "")
        # What does this data represent (e.g. stellar spectrum)
        meta["description"] = kwargs.pop("description", "")
        # Citation data (in bibtex format) if any
        meta["citation"] = kwargs.pop("citation", "")
        # Additional data for reference frame changes
        meta["star"] = kwargs.pop("star", None)
        meta["planet"] = kwargs.pop("planet", None)
        meta["observatory_location"] = kwargs.pop("observatory_location", None)
        # Datetime of the observation
        meta["datetime"] = kwargs.pop("datetime", Time(0, format="mjd"))
        # airmass is somewhat redundant (it can be calculated from star, observatory, and time)
        meta["airmass"] = kwargs.pop("airmass", None)
        # One of "barycentric", "telescope", "planet", "star"
        reference_frame = kwargs.pop("reference_frame", "barycentric")

        # Obsolete keywords
        kwargs.pop("sky_location", None)

        marked_for_death = []
        for key, value in kwargs.items():
            if key not in self.__init_args__:
                meta[key] = value
                marked_for_death += [key]

        for key in marked_for_death:
            del kwargs[key]

        kwmeta = kwargs.pop("meta", {})
        kwmeta.update(meta)
        kwargs["meta"] = kwmeta

        super().__init__(*args, **kwargs)

        self.reference_frame = reference_frame

    def __len__(self):
        return len(self.spectral_axis)

    def __copy__(self):
        cls = self.__class__
        wave = np.copy(self.wavelength)
        flux = np.copy(self.flux)
        meta = self.meta

        other = cls(flux=flux, spectral_axis=wave, **meta)
        return other

    def __deepcopy__(self, *args, **kwargs):
        cls = self.__class__
        wave = np.copy(self.wavelength)
        flux = np.copy(self.flux)
        meta = deepcopy(self.meta)

        other = cls(flux=flux, spectral_axis=wave, **meta)
        return other

    @property
    def regions(self):
        wmin, wmax = self.wavelength[[0, -1]]
        return specutils.SpectralRegion(wmin, wmax)

    def _get_header(self):
        header = {
            "SOURCE": self.meta["source"],
            "DESCR": self.meta["description"],
            "CITATION": self.meta["citation"].replace("\n", "").strip(),
            "DATE-OBS": self.datetime.fits,
            "DATE": datetime.now().isoformat(),
            "REFFRAME": str(self.reference_frame),
        }

        # This only saves the name
        # Values have to be recovered using various sources
        # StellarDB for star and planet
        # EarthCoordinates for observatory
        if self.meta["star"] is not None:
            header["STAR"] = str(self.meta["star"])
            header["RA"] = self.meta["star"].coordinates.ra.to_value("hourangle")
            header["DEC"] = self.meta["star"].coordinates.dec.to_value("deg")
        if self.meta["planet"] is not None:
            header["PLANET"] = str(self.meta["planet"])
        if self.meta["observatory_location"] is not None:
            if isinstance(self.meta["observatory_location"], str):
                header["OBSNAME"] = self.meta["observatory_location"]
            elif isinstance(self.meta["observatory_location"], (tuple, list)):
                header["OBSLAT"] = (self.meta["observatory_location"][0],)
                header["OBSLON"] = (self.meta["observatory_location"][1],)
                header["OBSALT"] = (self.meta["observatory_location"][2],)
        return header

    def get_fits_hdu(self):
        wave = self.wavelength.to(u.AA)
        flux = self.flux.decompose()
        wave_unit = str(wave.unit)
        flux_unit = str(flux.unit)
        wave = fits.Column(
            name="wavelength", array=wave.value, format="D", coord_unit=wave_unit
        )

        if flux.ndim == 1:
            flux = fits.Column(
                name="flux", array=flux.value, format="D", coord_unit=flux_unit,
            )
        elif flux.ndim == 2:
            flux = fits.Column(
                name="flux",
                array=flux.value.T,
                format=f"{flux.shape[0]}D",
                coord_unit=flux_unit,
            )
        else:
            flux = fits.Column(
                name="flux",
                array=flux.value.T,
                format=f"{flux.shape[0]}D",
                dim=flux.shape[1:],
                coord_unit=flux_unit,
            )

        header = self._get_header()
        header = fits.Header(header)
        hdu = fits.BinTableHDU.from_columns([wave, flux], header=header)
        return hdu

    @classmethod
    def read_fits_hdu(cls, hdu):
        header = hdu.header
        data = hdu.data
        meta = {}

        wunit = data.columns["wavelength"].coord_unit
        wunit = u.Unit(wunit) if wunit is not None else u.one
        funit = data.columns["flux"].coord_unit
        funit = u.Unit(funit) if funit is not None else u.one
        wave = data["wavelength"] << wunit
        flux = data["flux"] << funit

        if flux.ndim >= 2:
            flux = flux.T

        meta["source"] = header["SOURCE"]
        meta["description"] = header["DESCR"]
        meta["citation"] = header["CITATION"]
        meta["datetime"] = Time(header["DATE-OBS"], format="fits")
        meta["reference_frame"] = header["REFFRAME"]

        sdb = StellarDb()
        star = None
        if "STAR" in header:
            star = header["STAR"]
            try:
                star = sdb.get(star)
                meta["star"] = star
            except AttributeError:
                # TODO: warning, or empty star or something
                star = None
                pass
        if "PLANET" in header:
            planet = header["PLANET"]
            if star is not None:
                planet = star.planets[planet]
                meta["planet"] = planet
        if "OBSNAME" in header:
            meta["observatory_location"] = header["OBSNAME"]
        elif "OBSLON" in header and "OBSLAT" in header and "OBSALT" in header:
            meta["observatory_location"] = (
                header["OBSLON"],
                header["OBSLAT"],
                header["OBSALT"],
            )
        if "RA" in header and "DEC" in header:
            ra = coords.Angle(header["RA"], "hourangle")
            dec = header["DEC"] * u.deg
            if "star" in meta.keys():
                meta["star"].coordinates = coords.SkyCoord(ra, dec)

        spec = cls(flux=flux, spectral_axis=wave, **meta)

        return spec

    def shift(self, target_frame, inplace=False, rv=None):
        """
        Shift the spectrum from the current
        to a new reference frame

        Acceptable reference frames are:
          - barycentric
          - telescope (NOT barycentric)
          - planet
          - star

        Parameters
        ----------
        target_frame : str
            the NEW reference frame
        inplace : bool, optional
            whether to perform the data shift in place
            (overriding existing data), or not (default: False)

        Returns
        -------
        spec: Spectrum1D
            the shifted spectrum structure
        """

        # TODO: optimize in case that the target frame == reference frame
        # TODO: so we don't have to waste time on that

        # TODO: conversion to/from star/planet requires info about the star/planet
        # TODO: don't forget to use the radial velocity as well (?)
        try:
            target_frame = self.reference_frame_from_name(target_frame)
        except ValueError:
            pass

        # Step 1: determine relative velocity between current and target frame
        if rv is None:
            rv = self.reference_frame.to_frame(target_frame, self.datetime)

        # Step 2: Use the determined radial velocity to calculate a new wavelength grid
        beta = (rv / const.c).to_value(1)
        factor = np.sqrt((1 + beta) / (1 - beta))
        shifted = self.wavelength * factor

        # Step 3: Create new Spectrum1D with shifted wavelength grid
        if inplace:
            self._spectral_axis = shifted
            self.reference_frame = target_frame
            spec = self
        else:
            spec = Spectrum1D(spectral_axis=shifted, flux=self.flux, **self.meta)
            spec.reference_frame = target_frame

        return spec

    def resample(self, grid, method="linear", inplace=False, **kwargs):
        """
        Resample the current spectrum to a different wavelength grid

        Parameters
        ----------
        grid : Quantity
            The new wavelength grid
        method : str, optional
            The method to use for interpolation. Must be one of "flux_conserving", "linear", or "spline".
            By default "flux_conserving"
        ** kwargs
            Any optional keywords are passed on to the resampler. See specutils.manipulation for details

        Returns
        -------
        spec : Spectrum1D
            The resampled spectrum
        """
        options = ["flux_conserving", "linear", "spline"]
        if method == "flux_conserving":
            resampler = specman.FluxConservingResampler(**kwargs)
        elif method == "linear":
            resampler = specman.LinearInterpolatedResampler(**kwargs)
        elif method == "spline":
            resampler = specman.SplineInterpolatedResampler(**kwargs)
            self._data = np.nan_to_num(self._data)
        else:
            raise ValueError(
                f"Interpolation method not understood. Expected one of {options}, but got {method}"
            )

        spec = resampler(self, grid)

        if inplace:
            self._spectral_axis[:] = spec._spectral_axis
            self._data[:] = spec._data
            self._unit = spec._unit
            if self._uncertainty is not None:
                self._uncertainty = spec._uncertainty
            else:
                self._uncertainty = spec._uncertainty
        else:
            # Cast spec to Spectrum 1D class and set meta parameters
            spec.meta = copy(self.meta)
            spec.__class__ = Spectrum1D
            # spec.with_velocity_convention(self.velocity_convention)
            # spec.radial_velocity = self.radial_velocity

        return spec

    def extract_region(self, wrange):
        wave, flux = [], []

        if hasattr(wrange, "subregions"):
            wrange = wrange.subregions

        for wr in wrange:
            if hasattr(wr, "subregions"):
                wr = wr.subregions[0]
            wmin, wmax = wr

            mask = (self.wavelength >= wmin) & (self.wavelength <= wmax)

            wave += [self.wavelength[mask]]
            flux += [self.flux[..., mask]]

        if len(wrange) == 1:
            spec = Spectrum1D(flux=flux[0], spectral_axis=wave[0], **self.meta)
        else:
            segments = np.zeros(len(wrange) + 1)
            segments[1:] = np.cumsum([len(w) for w in wave])
            flux = np.hstack(flux)
            wave = np.concatenate(wave)
            wave = np.tile(wave, (flux.shape[0], 1))
            spec = SpectrumArray(
                flux=flux, spectral_axis=wave, segments=segments, **self.meta
            )

        return spec

class SpectrumList(Sequence, SpectrumBase):
    """
    Stores a list of Spectrum1D objects, with shared metadata
    This usually represents the different orders of the spectrum,
    which may have various sizes of spectral axis, especially when
    using model data

    Really this is only a thin shell with convenience functions,
    to act on all spectra at once

    NOTE
    ----
    No checks whatsoever are performed on the input data. Please
    only use CATS Spectra1D objects as input

    """

    def __init__(self, flux, spectral_axis, uncertainty=None, **kwargs):
        super().__init__()

        # We actually just pass everything to each individual spectrum
        # instead of trying to organize the metadata into one place
        # This means we don't have to worry about setting the values later
        # But this also means that they could change if we are not careful
        self._data = []
        if uncertainty is None:
            uncertainty = [None for _ in flux]
        for f, sa, unc in zip(flux, spectral_axis, uncertainty):
            spec = Spectrum1D(flux=f, spectral_axis=sa, uncertainty=unc, **kwargs)
            self._data += [spec]

    def __getitem__(self, key):
        return self._data[key]

    def __setitem__(self, key, value):
        # TODO: check that its a Spectrum1D ?
        self._data[key] = value

    def __len__(self):
        return len(self._data)

    def __operator__(self, other, operator):
        if isinstance(other, (float, int)) or (
            hasattr(other, "size") and other.size == 1
        ):
            # If its scalar, make it an array of the same length
            other = [other for _ in self]
        elif len(other) != len(self):
            raise ValueError(f"Incompatible sizes of {len(self)} and {len(other)}")
        # TODO: check that reference frame is the same
        data = [operator(t, o) for t, o in zip(self, other)]
        sl = self.__class__.from_spectra(data)
        for out, inp in zip(sl, self):
            out.meta = inp.meta
        return sl

    def __add__(self, other):
        return self.__operator__(other, op.add)

    def __sub__(self, other):
        return self.__operator__(other, op.sub)

    def __mul__(self, other):
        return self.__operator__(other, op.mul)

    def __truediv__(self, other):
        return self.__operator__(other, op.truediv)

    def __copy__(self):
        return SpectrumList.from_spectra(self._data)

    def __deepcopy__(self, *args, **kwargs):
        data = [deepcopy(s, *args, **kwargs) for s in self]
        return SpectrumList.from_spectra(data)

    @property
    def shape(self):
        return (len(self), [len(d) for d in self])

    @property
    def size(self):
        return sum(self.shape[1])

    @property
    def regions(self):
        wrange = [spec.regions for spec in self]
        for wr in wrange[1:]:
            wrange[0] += wr
        return wrange[0]

    @property
    def flux(self):
        return [s.flux for s in self]

    @property
    def wavelength(self):
        return [s.wavelength for s in self]

    @property
    def uncertainty(self):
        return [s.uncertainty for s in self]

    @property
    def meta(self):
        return self[0].meta

    @property
    def datetime(self):
        return self[0].datetime

    @datetime.setter
    def datetime(self, value):
        for s in self:
            s.datetime = value

    @property
    def reference_frame(self):
        return self[0].reference_frame

    @reference_frame.setter
    def reference_frame(self, value):
        for s in self:
            s.reference_frame = value

    @classmethod
    def from_spectra(cls, spectra):
        """
        Create a new SpectrumList object from a list of Spectrum1D

        Note
        ----
        The input spectra are NOT copied, therefore any changes
        made will be present in the existing variables

        Parameters
        ----------
        spectra : list
            list of Spectrum1D

        Returns
        -------
        spectrum_list : SpectrumList
            SpectrumList with containing the spectra
        """
        spectrum_list = cls([], [])
        spectrum_list._data = spectra
        return spectrum_list

    def write(self, fname, detector=None):
        """Save all spectra in SpectrumList to a single fits file
        Each Spectrum1D has its own extension with a complete header

        Parameters
        ----------
        fname : str
            output filename
        detector : Detector, optional
            Detector used, will add some data to the primary header. By default None
        """
        # Save all in one fits file
        # with metatdata in the header and the wavelength and flux in the data

        header = {}
        if detector is not None:
            header["INSTRUME"] = str(detector)
            header["INS SET"] = detector.setting

        header = fits.Header(header)
        primary = fits.PrimaryHDU(header=header)
        secondary = []
        for spec in self:
            secondary += [spec._get_fits_hdu()]

        hdulist = fits.HDUList(hdus=[primary, *secondary])
        hdulist.writeto(fname, overwrite=True)

    @classmethod
    def read(cls, fname):
        hdulist = fits.open(fname)
        header = hdulist[0].header

        det = None
        if "INSTRUME" in header:
            det = header["INSTRUME"]
            setting = header["INS SET"]
            det = det.capitalize().replace("+", "")

            module = getattr(detector, det)
            det = module(setting)

        spectra = []
        for i in range(1, len(hdulist)):
            spectra += [Spectrum1D._read_fits_hdu(hdulist[i])]
            if det is not None:
                spectra[-1].meta["observatory_location"] = det.observatory
                spectra[-1].reference_frame.observatory = det.observatory

        speclist = cls.from_spectra(spectra)
        return speclist

    def resample(self, grid, **kwargs):
        """
        Resample the different spectra onto the given grid

        Parameters
        ----------
        grid : list
            list of wavelengths arrays to resample to
        ** kwargs
            keyword arguments passed to Spectrum1D resample

        Returns
        -------
        SpectrumList
            New resampled SpectrumList object
        """

        if len(grid) != len(self):
            raise ValueError(
                f"Wavelength grid has {len(grid)} arrays, but SpectrumList has {len(self)} spectra."
            )

        data = []
        for i, (g, spec) in enumerate(zip(grid, self)):
            s = spec.resample(g, **kwargs)
            data += [s]

        sl = SpectrumList.from_spectra(data)
        return sl

    def shift(self, target_frame, inplace=False, rv=None, **kwargs):
        """
        Shift all spectra to the target frame

        Parameters
        ----------
        target_frame : ReferenceFrame
            target reference frame to convert the Spectra to
        ** kwargs
            keyword arguments passed to Spectrum1D shift
        Returns
        -------
        SpectrumList
            New shifted SpectrumList object
        """

        try:
            target_frame = self[0].reference_frame_from_name(target_frame)
        except ValueError:
            pass

        if rv is None:
            rv = self.reference_frame.to_frame(target_frame, self.datetime)

        data = []
        for spec in self:
            s = spec.shift(target_frame, inplace=inplace, rv=rv, **kwargs)
            data += [s]

        if not inplace:
            sl = SpectrumList.from_spectra(data)
            return sl
        else:
            return self

class SpectrumArray(SpectrumBase, Sequence):
    """
    A Collection of Spectra with the same size,
    but possibly different wavelength axis
    """

    def __init__(self, *args, **kwargs):
        if len(args) == 1:
            if isinstance(args[0], list):
                spectra = args[0]
            elif isinstance(args[0], SpectrumList):
                spectra = [args[0]]
            else:
                raise ValueError

            nspec = len(spectra)
            nseg = spectra[0].shape[0]
            npix = spectra[0].size

            wunit = spectra[0][0].wavelength.unit
            funit = spectra[0][0].flux.unit

            self.wavelength = np.zeros((nspec, npix)) << wunit
            self.flux = np.zeros((nspec, npix)) << funit
            for i, spec in enumerate(spectra):
                self.wavelength[i] = np.concatenate(spec.wavelength)
                self.flux[i] = np.concatenate(spec.flux)

            if spectra[0][0].uncertainty is not None:
                uunit = spectra[0][0].uncertainty.unit
                self.uncertainty = np.zeros((nspec, npix)) << uunit
                for i, spec in enumerate(spectra):
                    self.uncertainty[i] = np.concatenate(
                        [unc.array for unc in spec.uncertainty]
                    )
                self.uncertainty = StdDevUncertainty(self.uncertainty)
            else:
                self.uncertainty = None

            self.meta = {}
            self.segments = np.zeros(nseg + 1, dtype=int)
            self.segments[1:] = spectra[0].shape[1]
            self.segments = np.cumsum(self.segments)

            self.meta.update(spectra[0].meta)
            times = Time([spec.datetime for spec in spectra])
            self.meta["datetime"] = times
        elif len(args) == 0 and len(kwargs) > 0:
            self.meta = {}
            self.wavelength = kwargs.pop("spectral_axis")
            self.flux = kwargs.pop("flux")
            self.segments = kwargs.pop("segments")
            self.uncertainty = kwargs.pop("uncertainty", None)
            self.meta.update(kwargs)
        else:
            raise ValueError

    def __len__(self):
        return len(self.wavelength)

    @property
    def nseg(self):
        return len(self.segments) - 1

    @property
    def segments(self):
        return self.meta["segments"]

    @segments.setter
    def segments(self, value):
        self.meta["segments"] = value

    def __getitem__(self, key):
        wave = self.wavelength[key]
        flux = self.flux[key]
        if self.uncertainty is not None:
            uncs = self.uncertainty[key]
        else:
            uncs = None

        spectra = []
        for left, right in zip(self.segments[:-1], self.segments[1:]):
            meta = self.meta.copy()
            meta["datetime"] = self.datetime[key]
            if uncs is not None:
                meta["uncertainty"] = uncs[left:right]
            spec = Spectrum1D(
                flux=flux[left:right], spectral_axis=wave[left:right], **meta
            )
            spectra += [spec]

        speclist = SpectrumList.from_spectra(spectra)
        return speclist

    def __setitem__(self, key, value):
        self.wavelength[key] = np.concatenate(value.wavelength)
        self.flux[key] = np.concatenate(value.flux)

    def __operator__(self, other, operator):
        if isinstance(other, (float, int)) or (
            hasattr(other, "size") and other.size == 1
        ):
            data = operator(self.flux, other)
        elif isinstance(other, SpectrumArray):
            data = operator(self.flux, other.flux)
        elif isinstance(other, np.ndarray):
            data = operator(self.flux, other)
        else:
            return NotImplemented

        sa = self.__class__(
            flux=data,
            spectral_axis=self.wavelength,
            segments=self.segments,
            **self.meta,
        )
        return sa

    def __add__(self, other):
        return self.__operator__(other, op.add)

    def __sub__(self, other):
        return self.__operator__(other, op.sub)

    def __mul__(self, other):
        return self.__operator__(other, op.mul)

    def __truediv__(self, other):
        return self.__operator__(other, op.truediv)

    @property
    def shape(self):
        return self.wavelength.shape

    def get_segment(self, seg):
        left, right = self.segments[seg : seg + 2]
        left, right = int(left), int(right)
        wave = self.wavelength[:, left:right]
        flux = self.flux[:, left:right]
        if self.uncertainty is not None:
            uncs = self.uncertainty[:, left : right]
        else:
            uncs = None
        specarr = SpectrumArray(
            flux=flux,
            spectral_axis=wave,
            uncertainty=uncs,
            segments=[0, right - left],
            **{k: v for k, v in self.meta.items() if k != "segments"},
        )
        return specarr

    @classmethod
    def from_dict(cls, data):
        meta = data["meta"] if "meta" in data.keys() else {}
        flux_unit = data["flux_unit"]
        wave_unit = data["spectral_axis_unit"]
        flux = data["flux"] << u.Unit(flux_unit)
        wave = data["spectral_axis"] << u.Unit(wave_unit)
        segments = data["segments"]

        if "uncertainty" in data.keys():
            uncs_unit = data["uncertainty_unit"]
            uncs = data["uncertainty"] << u.Unit(uncs_unit)
            uncs = StdDevUncertainty(uncs)
            meta["uncertainty"] = uncs

        self = cls(flux=flux, spectral_axis=wave, segments=segments, **meta)
        return self

    def shift(self, target_frame, inplace=True, rv=None):
        if not inplace:
            spectra = deepcopy(self)
        else:
            spectra = self

        try:
            target_frame = self.reference_frame_from_name(target_frame)
        except ValueError:
            pass

        if rv is None:
            rv = self.reference_frame.to_frame(target_frame, self.datetime)

        meta = self.meta.copy()
        for i in tqdm(range(len(spectra)), leave=False):
            meta["datetime"] = spectra.datetime[i]
            if spectra.uncertainty is not None:
                meta["uncertainty"] = spectra.uncertainty[i]

            spec = Spectrum1D(
                flux=spectra.flux[i], spectral_axis=spectra.wavelength[i], **meta
            )
            spec = spec.shift(target_frame, rv=rv[i], inplace=True)
            spectra.wavelength[i] = spec.wavelength
            target_frame = spec.reference_frame

        spectra.meta["reference_frame"] = target_frame
        return spectra

    def resample(self, wavelength, inplace=False, **kwargs):
        n_spectra = self.flux.shape[0]

        if np.ndim(wavelength) == 1:
            wavelength = np.tile(wavelength, (n_spectra, 1))

        n_wave_points = wavelength.shape[1]

        if not inplace:
            spectra = deepcopy(self)
            spectra.wavelength = np.copy(wavelength)
            spectra.flux = (
                np.zeros((n_spectra, n_wave_points), dtype=spectra.flux.dtype)
                << spectra.flux.unit
            )
            spectra.segments = [0, len(wavelength)]
        else:
            spectra = self

        for i in tqdm(range(len(self)), leave=False):
            meta = {}
            if self.uncertainty is not None:
                meta["uncertainty"] = self.uncertainty[i]
            spec = Spectrum1D(
                flux=self.flux[i], spectral_axis=self.wavelength[i], **meta
            )
            spec = spec.resample(wavelength[i], inplace=inplace, **kwargs)
            if not inplace:
                spectra.flux[i] = spec.flux
            spectra.wavelength[i] = wavelength[i]

        return spectra

class SpectrumArrayIO(StepIO):
    def save(self, data, fname=None):
        if fname is None:
            fname = self.savefilename
        data.write(fname)

    def load(self, fname=None):
        if fname is None:
            fname = self.savefilename
        data = SpectrumArray.read(fname)
        return data

class Spectrum1DIO(StepIO):
    def save(self, data, fname=None):
        if fname is None:
            fname = self.savefilename
        data.write(fname)

    def load(self, fname=None):
        if fname is None:
            fname = self.savefilename
        data = Spectrum1D.read(fname)
        return data

import numpy,import astropy,from astropy


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

from collections import namedtuple

import numpy as np

Rule = namedtuple("Rule", ["dtype", "multiplier", "round"])

class Converter(object):
    def __init__(self, rules):
        self._rules = rules

    def convert(self, name, data):
        try:
            r = self._rules[name]
        except KeyError:
            return data

        result = data * r.multiplier
        if r.round:
            result = np.round(result, r.round)

        return result

    def field_type(self, name, dt):
        try:
            return self._rules[name].dtype
        except KeyError:
            return dt

float64 = np.dtype("float64")

StockBarConverter = Converter(
    {
        "open": Rule(float64, 1 / 10000.0, 2),
        "close": Rule(float64, 1 / 10000.0, 2),
        "high": Rule(float64, 1 / 10000.0, 2),
        "low": Rule(float64, 1 / 10000.0, 2),
        "limit_up": Rule(float64, 1 / 10000.0, 2),
        "limit_down": Rule(float64, 1 / 10000.0, 2),
        "volume": Rule(float64, 1, 0),
    }
)

FutureDayBarConverter = Converter(
    {
        "open": Rule(float64, 1 / 10000.0, 3),
        "close": Rule(float64, 1 / 10000.0, 3),
        "high": Rule(float64, 1 / 10000.0, 3),
        "low": Rule(float64, 1 / 10000.0, 3),
        "limit_up": Rule(float64, 1 / 10000.0, 3),
        "limit_down": Rule(float64, 1 / 10000.0, 3),
        "volume": Rule(float64, 1, 0),
        "basis_spread": Rule(float64, 1 / 10000.0, 4),
        "settlement": Rule(float64, 1 / 10000.0, 3),
        "prev_settlement": Rule(float64, 1 / 10000.0, 3),
    }
)

FundDayBarConverter = Converter(
    {
        "open": Rule(float64, 1 / 10000.0, 3),
        "close": Rule(float64, 1 / 10000.0, 3),
        "high": Rule(float64, 1 / 10000.0, 3),
        "low": Rule(float64, 1 / 10000.0, 3),
        "acc_net_value": Rule(float64, 1 / 10000.0, 4),
        "unit_net_value": Rule(float64, 1 / 10000.0, 4),
        "discount_rate": Rule(float64, 1 / 10000.0, 4),
        "limit_up": Rule(float64, 1 / 10000.0, 4),
        "limit_down": Rule(float64, 1 / 10000.0, 4),
        "volume": Rule(float64, 1, 0),
    }
)

IndexBarConverter = Converter(
    {
        "open": Rule(float64, 1 / 10000.0, 2),
        "close": Rule(float64, 1 / 10000.0, 2),
        "high": Rule(float64, 1 / 10000.0, 2),
        "low": Rule(float64, 1 / 10000.0, 2),
        "volume": Rule(float64, 1, 0),
    }
)

PublicFundDayBarConverter = Converter(
    {
        "open": Rule(float64, 1 / 10000.0, 4),
        "close": Rule(float64, 1 / 10000.0, 4),
        "high": Rule(float64, 1 / 10000.0, 4),
        "low": Rule(float64, 1 / 10000.0, 4),
        "acc_net_value": Rule(float64, 1 / 10000.0, 4),
        "unit_net_value": Rule(float64, 1 / 10000.0, 4),
        "limit_up": Rule(float64, 1 / 10000.0, 4),
        "limit_down": Rule(float64, 1 / 10000.0, 4),
    }
)


from collections import OrderedDict
import numpy as np

import torch
import torch.optim as optim
from torch import nn as nn

import rlkit.torch.pytorch_util as ptu
from rlkit.core.eval_util import create_stats_ordered_dict
from rlkit.core.rl_algorithm import MetaRLAlgorithm

import torch.nn.functional as F
from torch.optim import Adam
from sac.utils import soft_update,hard_update

class PEARLSoftActorCritic(MetaRLAlgorithm):
    def __init__(
            self,
            env,
            train_tasks,
            eval_tasks,
            latent_dim,
            nets,

            policy_lr=3e-4,
            qf_lr=1e-3,
            vf_lr=1e-3,
            alpha = 0.15,
            automatic_entropy_tuning = True,
            lr=3e-4,
            context_lr=3e-4,
            kl_lambda=1.,
            policy_mean_reg_weight=1e-3,
            policy_std_reg_weight=1e-3,
            policy_pre_activation_weight=0.,
            optimizer_class=optim.Adam,
            recurrent=False,
            use_information_bottleneck=True,
            sparse_rewards=False,

            #soft_target_tau=1e-2,
            soft_target_tau=0.005,
            plotter=None,
            render_eval_paths=False,
            **kwargs
    ):
        super().__init__(
            env=env,
            agent=nets[0],
            train_tasks=train_tasks,
            eval_tasks=eval_tasks,
            **kwargs
        )

        self.soft_target_tau = soft_target_tau
        self.policy_mean_reg_weight = policy_mean_reg_weight
        self.policy_std_reg_weight = policy_std_reg_weight
        self.policy_pre_activation_weight = policy_pre_activation_weight
        self.plotter = plotter
        self.render_eval_paths = render_eval_paths
        self.alpha = alpha
        self.automatic_entropy_tuning = automatic_entropy_tuning

        self.recurrent = recurrent
        self.latent_dim = latent_dim
        # self.qf_criterion = nn.MSELoss()
        # self.vf_criterion = nn.MSELoss()
        self.vib_criterion = nn.MSELoss()
        self.l2_reg_criterion = nn.MSELoss()
        self.kl_lambda = kl_lambda

        self.use_information_bottleneck = use_information_bottleneck
        self.sparse_rewards = sparse_rewards

        # self.qf1, self.qf2, self.vf = nets[1:]
        self.critic,self.critic_target = nets[1:]#q1,q2,target q1,target q2
        # self.target_vf = self.vf.copy()
        self.critic_optimizer = Adam(self.critic.parameters(), lr=lr)
        hard_update(self.critic_target,self.critic)
        self.policy_optimizer = Adam(self.agent.policy.parameters(), lr=policy_lr)

        self.target_entropy = -torch.prod(torch.Tensor(env.action_space.shape)).to("cuda").item()  # torch.prod(input) : 返回所有元素的乘积
        self.log_alpha = torch.zeros(1, requires_grad=True, device="cuda")
        self.alpha_optim = Adam([self.log_alpha], lr=lr)

        # self.policy_optimizer = optimizer_class(
        #     self.agent.policy.parameters(),
        #     lr=policy_lr,
        # )
        # self.qf1_optimizer = optimizer_class(
        #     self.qf1.parameters(),
        #     lr=qf_lr,
        # )
        # self.qf2_optimizer = optimizer_class(
        #     self.qf2.parameters(),
        #     lr=qf_lr,
        # )
        # self.vf_optimizer = optimizer_class(
        #     self.vf.parameters(),
        #     lr=vf_lr,
        # )
        self.context_optimizer = optimizer_class(
            self.agent.context_encoder.parameters(),
            lr=context_lr,
        )

    ###### Torch stuff #####
    @property
    def networks(self):
        return self.agent.networks + [self.agent] + [self.critic, self.critic_target]

    def training_mode(self, mode):
        for net in self.networks:
            net.train(mode)

    def to(self, device=None):
        if device == None:
            device = ptu.device
        for net in self.networks:
            net.to(device)

    ##### Data handling #####
    def sample_data(self, indices, encoder=False):
        ''' sample data from replay buffers to construct a training meta-batch '''
        # collect data from multiple tasks for the meta-batch
        obs, actions, rewards, next_obs, terms = [], [], [], [], []
        for idx in indices:
            if encoder:
                batch = ptu.np_to_pytorch_batch(self.enc_replay_buffer.random_batch(idx, batch_size=self.embedding_batch_size, sequence=self.recurrent))
            else:
                batch = ptu.np_to_pytorch_batch(self.replay_buffer.random_batch(idx, batch_size=self.batch_size))
            o = batch['observations'][None, ...]
            a = batch['actions'][None, ...]
            if encoder and self.sparse_rewards:
                # in sparse reward settings, only the encoder is trained with sparse reward
                r = batch['sparse_rewards'][None, ...]
            else:
                r = batch['rewards'][None, ...]
            no = batch['next_observations'][None, ...]
            t = batch['terminals'][None, ...]
            obs.append(o)
            actions.append(a)
            rewards.append(r)
            next_obs.append(no)
            terms.append(t)
        obs = torch.cat(obs, dim=0)
        actions = torch.cat(actions, dim=0)
        rewards = torch.cat(rewards, dim=0)
        next_obs = torch.cat(next_obs, dim=0)
        terms = torch.cat(terms, dim=0)
        return [obs, actions, rewards, next_obs, terms]

    def prepare_encoder_data(self, obs, act, rewards):
        ''' prepare context for encoding '''
        # for now we embed only observations and rewards
        # assume obs and rewards are (task, batch, feat)
        task_data = torch.cat([obs, act, rewards], dim=2)
        return task_data

    def prepare_context(self, idx):
        ''' sample context from replay buffer and prepare it '''
        batch = ptu.np_to_pytorch_batch(self.enc_replay_buffer.random_batch(idx, batch_size=self.embedding_batch_size, sequence=self.recurrent))
        obs = batch['observations'][None, ...]
        act = batch['actions'][None, ...]
        rewards = batch['rewards'][None, ...]
        context = self.prepare_encoder_data(obs, act, rewards)
        return context

    ##### Training #####
    def _do_training(self, indices):#对indices中的任务进行训练
        mb_size = self.embedding_mini_batch_size
        num_updates = self.embedding_batch_size // mb_size

        batch = self.sample_data(indices, encoder=True)

        # zero out context and hidden encoder state
        self.agent.clear_z(num_tasks=len(indices))

        for i in range(num_updates):
            mini_batch = [x[:, i * mb_size: i * mb_size + mb_size, :] for x in batch]
            obs_enc, act_enc, rewards_enc, _, _ = mini_batch
            context = self.prepare_encoder_data(obs_enc, act_enc, rewards_enc)
            self._take_step(indices, context)

            # stop backprop
            self.agent.detach_z()

    def _take_step(self, indices, context):

        num_tasks = len(indices)

        # data is (task, batch, feat)
        obs, actions, rewards, next_obs, terms = self.sample_data(indices)#从Replay Buffer采集数据，s,a,r,s',d
        # obs, actions, rewards, next_obs, terms = self.sample_sac(indices)

        # run inference in networks
        policy_outputs, task_z = self.agent(obs, context)  # 策略forward的输出，以及任务隐变量Z
        new_actions, policy_mean, policy_log_std, log_pi = policy_outputs[:4]  # 下一个状态下策略所采取的动作，其log概率 line 63

        # flattens out the task dimension:
        t, b, _ = obs.size()
        obs = obs.view(t * b, -1)
        actions = actions.view(t * b, -1)
        next_obs = next_obs.view(t * b, -1)

        # Q networks
        # encoder will only get gradients from Q nets
        q1, q2 = self.critic(obs, actions, task_z)  # forward

        # get targets for use in V and Q updates
        with torch.no_grad():
            q1_next_target, q2_next_target = self.critic_target(next_obs, new_actions,task_z)  # target q               line 64
            min_qf_next_target = torch.min(q1_next_target, q2_next_target)  # 计算较小的target Q         line 65

        # KL constraint on z if probabilistic
        self.context_optimizer.zero_grad()
        if self.use_information_bottleneck:
            kl_div = self.agent.compute_kl_div()  # context_encoder前向传播得到μ和σ，计算该分布与先验分布的kl散度
            kl_loss = self.kl_lambda * kl_div
            kl_loss.backward(retain_graph=True)

        # qf and encoder update (note encoder does not get grads from policy or vf)
        self.critic_optimizer.zero_grad()
        rewards_flat = rewards.view(self.batch_size * num_tasks, -1)
        # scale rewards for Bellman update
        rewards_flat = rewards_flat * self.reward_scale
        terms_flat = terms.view(self.batch_size * num_tasks, -1)
        q_target = rewards_flat + (1. - terms_flat) * self.discount * (min_qf_next_target - self.alpha * log_pi)  # q=r+(1-d)γ(Vst+1)
        qf_loss = torch.mean((q1 - q_target) ** 2) + torch.mean((q2 - q_target) ** 2)
        qf_loss.backward()
        self.critic_optimizer.step()
        self.context_optimizer.step()

        # q1_loss = F.mse_loss(q1,q_target)  # JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]       line 69
        # q2_loss = F.mse_loss(q2,q_target)  # line 70
        # pi, log_pi, _ = self.agent.policy.sample(obs)  # 动作，动作的对数概率
        # print(obs.size())  #[1024,27]
        # print(task_z.size())

        # compute min Q on the new actions
        # in_policy = torch.cat([obs, task_z], 1)
        # pi, _, _, log_pi, _, _, _, _, = self.agent.policy(in_policy)  # line 72
        q1_pi, q2_pi = self.critic(obs, new_actions,task_z)  # 动作的Q值                                                line 74
        min_q_pi = torch.min(q1_pi, q2_pi)  # line 75

        # self.critic_optimizer.zero_grad()
        # q1_loss.backward(retain_graph=True)
        # self.critic_optimizer.step()

        # self.critic_optimizer.zero_grad()
        # q2_loss.backward(retain_graph=True)
        # self.critic_optimizer.step()

        #update target Q network
        soft_update(self.critic_target, self.critic, self.soft_target_tau)

        policy_loss = ((self.alpha * log_pi) - min_q_pi).mean()  # line 77
        mean_reg_loss = self.policy_mean_reg_weight * (policy_mean ** 2).mean()
        std_reg_loss = self.policy_std_reg_weight * (policy_log_std ** 2).mean()
        pre_tanh_value = policy_outputs[-1]
        pre_activation_reg_loss = self.policy_pre_activation_weight * (
            (pre_tanh_value ** 2).sum(dim=1).mean()
        )
        policy_reg_loss = mean_reg_loss + std_reg_loss + pre_activation_reg_loss
        policy_loss = policy_loss + policy_reg_loss

        self.policy_optimizer.zero_grad()
        policy_loss.backward(retain_graph=True)
        self.policy_optimizer.step()

        if self.automatic_entropy_tuning:
            alpha_loss = -(self.log_alpha * (log_pi + self.target_entropy).detach()).mean()  # E[-αlogπ(at|st)-αH]
            self.alpha_optim.zero_grad()
            alpha_loss.backward(retain_graph=True)
            self.alpha_optim.step()
            self.alpha = self.log_alpha.exp()
        else:
            alpha_loss = 0
            alpha = 1

        # save some statistics for eval
        if self.eval_statistics is None:
            # eval should set this to None.
            # this way, these statistics are only computed for one batch.
            self.eval_statistics = OrderedDict()
            if self.use_information_bottleneck:
                z_mean = np.mean(np.abs(ptu.get_numpy(self.agent.z_means[0])))
                z_sig = np.mean(ptu.get_numpy(self.agent.z_vars[0]))
                self.eval_statistics['Z mean train'] = z_mean
                self.eval_statistics['Z variance train'] = z_sig
                self.eval_statistics['KL Divergence'] = ptu.get_numpy(kl_div)
                self.eval_statistics['KL Loss'] = ptu.get_numpy(kl_loss)

            self.eval_statistics['Q1 Loss'] = np.mean(ptu.get_numpy(q1_loss))
            self.eval_statistics['Q2 Loss'] = np.mean(ptu.get_numpy(q2_loss))
            self.eval_statistics['Policy Loss'] = np.mean(ptu.get_numpy(
                policy_loss
            ))
            if self.automatic_entropy_tuning:
                self.eval_statistics['Alpha Loss'] = np.mean(ptu.get_numpy(alpha_loss))
                self.eval_statistics['Alpha'] = np.mean(ptu.get_numpy(self.alpha))
            self.eval_statistics.update(create_stats_ordered_dict(
                'Q1 Predictions',
                ptu.get_numpy(q1),
            ))
            self.eval_statistics.update(create_stats_ordered_dict(
                'Q2 Predictions',
                ptu.get_numpy(q2),
            ))
            self.eval_statistics.update(create_stats_ordered_dict(
                'Log Pis',
                ptu.get_numpy(log_pi),
            ))
            self.eval_statistics.update(create_stats_ordered_dict(
                'Policy mu',
                ptu.get_numpy(policy_mean),
            ))
            self.eval_statistics.update(create_stats_ordered_dict(
                'Policy log std',
                ptu.get_numpy(policy_log_std),
            ))

    def get_epoch_snapshot(self, epoch):
        # NOTE: overriding parent method which also optionally saves the env
        snapshot = OrderedDict(
            # qf1=self.qf1.state_dict(),
            # qf2=self.qf2.state_dict(),
            critic = self.critic.state_dict(),
            critic_target = self.critic_target.state_dict(),
            policy=self.agent.policy.state_dict(),
            # vf=self.vf.state_dict(),
            # target_vf=self.target_vf.state_dict(),
            context_encoder=self.agent.context_encoder.state_dict(),
        )
        return snapshot

rlkit/torch/sac/remain policy reg loss/sac.py


import hydra
import torch
from utils.util import AverageMeter, convert_layers
from utils.dataset import BlenderLaptopAuxDataset
import numpy as np
import logging
import torch.optim as optim
import torch.nn as nn
from tqdm import tqdm
from torchvision.models import segmentation
import os

@hydra.main(config_path='./config', config_name='laptop_aux')
def main(cfg):
    logger = logging.getLogger(__name__)

    # load name
    name_path = hydra.utils.to_absolute_path('data/shapenet_names/{}.txt'.format(cfg.category))
    if os.path.exists(name_path):
        shapenames = open(name_path).read().splitlines()

    ds = BlenderLaptopAuxDataset(cfg, shapenames)
    df = torch.utils.data.DataLoader(ds, pin_memory=True, batch_size=cfg.batch_size, shuffle=True, num_workers=10)

    segmenter = segmentation.fcn_resnet50(num_classes=2).cuda()
    segmenter = convert_layers(segmenter, nn.BatchNorm2d, nn.InstanceNorm2d).cuda()

    opt = optim.Adam([*segmenter.parameters()], lr=cfg.opt.lr, weight_decay=cfg.opt.weight_decay)
    ce = nn.CrossEntropyLoss()

    logger.info('Train')
    for epoch in range(cfg.max_epoch):
        n = 0

        loss_meter = AverageMeter()
        segmenter.train()
        with tqdm(df) as t:
            for rgb, label in t:
                rgb, label = rgb.cuda(), label.cuda()

                opt.zero_grad()

                feat = segmenter(rgb.permute(0, 3, 1, 2))['out']  # B x 2 x H x W

                loss = ce(feat, label)
                loss.backward(retain_graph=False)

                opt.step()

                loss_meter.update(loss.item())

                n += 1
                t.set_postfix(loss=loss_meter.avg)

        torch.save(segmenter.state_dict(), f'segmenter_current.pth')
        logger.info('loss: {:.4f}'.format(loss_meter.avg))

if __name__ == '__main__':
    main()


# Copyright 2021 DeepMind Technologies Limited.
#
# 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.

"""Script to generate ensembled PCQ test predictions."""

import collections
import os
import pathlib
from typing import List, NamedTuple

from absl import app
from absl import flags
from absl import logging
import dill
import numpy as np
from ogb import lsc

# pylint: disable=g-bad-import-order
# pytype: disable=import-error
import datasets

_NUM_SEEDS = 2

_CLIP_VALUE = 20.

_NUM_KFOLD_SPLITS = 10

_SEED_START = flags.DEFINE_integer(
    'seed_start', 42, 'Initial seed for the list of ensemble models.')

_CONFORMER_PATH = flags.DEFINE_string(
    'conformer_path', None, 'Path to conformer predictions.', required=True)

_NON_CONFORMER_PATH = flags.DEFINE_string(
    'non_conformer_path',
    None,
    'Path to non-conformer predictions.',
    required=True)

_OUTPUT_PATH = flags.DEFINE_string('output_path', None, 'Output path.')

_SPLIT = flags.DEFINE_enum('split', 'test', ['test', 'valid'],
                           'Split: valid or test.')

class _Predictions(NamedTuple):
  predictions: np.ndarray
  indices: np.ndarray

def _load_dill(fname) -> bytes:
  with open(fname, 'rb') as f:
    return dill.load(f)

def _sort_by_indices(predictions: _Predictions) -> _Predictions:
  order = np.argsort(predictions.indices)
  return _Predictions(
      predictions=predictions.predictions[order],
      indices=predictions.indices[order])

def load_predictions(path: str, split: str) -> _Predictions:
  """Load written prediction file."""
  if len(os.listdir(path)) != 1:
    raise ValueError('Prediction directory must have exactly '
                     'one prediction sub-directory: %s' % path)

  prediction_subdir = os.listdir(path)[0]
  return _Predictions(*_load_dill(f'{path}/{prediction_subdir}/{split}.dill'))

def mean_mae_distance(x, y):
  return np.abs(x - y).mean()

def _load_valid_labels() -> np.ndarray:
  labels = [label for _, label in datasets.load_smile_strings(with_labels=True)]
  return np.array([labels[i] for i in datasets.load_splits()['valid']])

def evaluate_valid_predictions(ensembled_predictions: _Predictions):
  """Evaluates the predictions on the validation set."""
  ensembled_predictions = _sort_by_indices(ensembled_predictions)
  evaluator = lsc.PCQM4MEvaluator()
  results = evaluator.eval(
      dict(
          y_pred=ensembled_predictions.predictions,
          y_true=_load_valid_labels()))
  logging.info('MAE on validation dataset: %f', results['mae'])

def clip_predictions(predictions: _Predictions) -> _Predictions:
  return predictions._replace(
      predictions=np.clip(predictions.predictions, 0., _CLIP_VALUE))

def _generate_test_prediction_file(test_predictions: np.ndarray,
                                   output_path: pathlib.Path) -> pathlib.Path:
  """Generates the final file for submission."""

  # Check that predictions are not nuts.
  assert test_predictions.dtype in [np.float64, np.float32]
  assert not np.any(np.isnan(test_predictions))
  assert np.all(np.isfinite(test_predictions))
  assert test_predictions.min() >= 0.
  assert test_predictions.max() <= 40.

  # Too risky to overwrite.
  if output_path.exists():
    raise ValueError(f'{output_path} already exists')

  # Write to a local directory, and copy to final path (possibly cns).
  # It is not possible to write directlt on CNS.
  evaluator = lsc.PCQM4MEvaluator()

  evaluator.save_test_submission(dict(y_pred=test_predictions), output_path)
  return output_path

def merge_complementary_results(split: str, results_a: _Predictions,
                                results_b: _Predictions) -> _Predictions:
  """Merges two prediction results with no overlap."""

  indices_a = set(results_a.indices)
  indices_b = set(results_b.indices)
  assert not indices_a.intersection(indices_b)

  if split == 'test':
    merged_indices = list(sorted(indices_a | indices_b))
    expected_indices = datasets.load_splits()[split]
    assert np.all(expected_indices == merged_indices)

  predictions = np.concatenate([results_a.predictions, results_b.predictions])
  indices = np.concatenate([results_a.indices, results_b.indices])
  predictions = _sort_by_indices(
      _Predictions(indices=indices, predictions=predictions))
  return predictions

def ensemble_valid_predictions(
    predictions_list: List[_Predictions]) -> _Predictions:
  """Ensembles a list of predictions."""
  index_to_predictions = collections.defaultdict(list)
  for predictions in predictions_list:
    for idx, pred in zip(predictions.indices, predictions.predictions):
      index_to_predictions[idx].append(pred)

  for idx, ensemble_list in index_to_predictions.items():
    if len(ensemble_list) != _NUM_SEEDS:
      raise RuntimeError(
          'Graph index in the validation set received wrong number of '
          'predictions to ensemble.')

  index_to_predictions = {
      k: np.median(pred_list, axis=0)
      for k, pred_list in index_to_predictions.items()
  }
  return _sort_by_indices(
      _Predictions(
          indices=np.array(list(index_to_predictions.keys())),
          predictions=np.array(list(index_to_predictions.values()))))

def ensemble_test_predictions(
    predictions_list: List[_Predictions]) -> _Predictions:
  """Ensembles a list of predictions."""
  predictions = np.median([pred.predictions for pred in predictions_list],
                          axis=0)
  common_indices = predictions_list[0].indices
  for preds in predictions_list[1:]:
    assert np.all(preds.indices == common_indices)

  return _Predictions(predictions=predictions, indices=common_indices)

def create_submission_from_predictions(
    output_path: pathlib.Path, test_predictions: _Predictions) -> pathlib.Path:
  """Creates a submission for predictions on a path."""
  assert _SPLIT.value == 'test'

  output_path = _generate_test_prediction_file(
      test_predictions.predictions,
      output_path=output_path / 'submission_files')

  return output_path / 'y_pred_pcqm4m.npz'

def merge_predictions(split: str) -> List[_Predictions]:
  """Generates features merged from conformer and non-conformer predictions."""
  merged_predictions: List[_Predictions] = []
  seed = _SEED_START.value

  # Load conformer and non-conformer predictions.
  for unused_seed_group in (0, 1):
    for k in range(_NUM_KFOLD_SPLITS):
      conformer_predictions: _Predictions = load_predictions(
          f'{_CONFORMER_PATH.value}/k{k}_seed{seed}', split)

      non_conformer_predictions: _Predictions = load_predictions(
          f'{_NON_CONFORMER_PATH.value}/k{k}_seed{seed}', split)

      merged_predictions.append(
          merge_complementary_results(_SPLIT.value, conformer_predictions,
                                      non_conformer_predictions))

      seed += 1
  return merged_predictions

def main(_):
  split: str = _SPLIT.value

  # Merge conformer and non-conformer predictions.
  merged_predictions = merge_predictions(split)

  # Clip before ensembling.
  clipped_predictions = list(map(clip_predictions, merged_predictions))

  # Ensemble predictions.
  if split == 'valid':
    ensembled_predictions = ensemble_valid_predictions(clipped_predictions)
  else:
    assert split == 'test'
    ensembled_predictions = ensemble_test_predictions(clipped_predictions)

  # Clip after ensembling.
  ensembled_predictions = clip_predictions(ensembled_predictions)

  ensembled_predictions_path = pathlib.Path(_OUTPUT_PATH.value)
  ensembled_predictions_path.mkdir(parents=True, exist_ok=True)

  with open(ensembled_predictions_path / f'{split}_predictions.dill',
            'wb') as f:
    dill.dump(ensembled_predictions, f)

  if split == 'valid':
    evaluate_valid_predictions(ensembled_predictions)
  else:
    assert split == 'test'
    output_path = create_submission_from_predictions(ensembled_predictions_path,
                                                     ensembled_predictions)
    logging.info('Submission files written to %s', output_path)

if __name__ == '__main__':
  app.run(main)


#!/usr/bin/env python3
import os
import base
import numpy as np
def tasks():
    return [
        MutationDimosLocations(),
        ]
class MutationDimosLocations(base.Task):
    """
    Retrieves the location of mutations in SCN5A, based on the homology
    modeling by Dimos.
    """
    def __init__(self):
        super(MutationDimosLocations, self).__init__(
            'mutation_dimos_locations')
        self._set_data_subdir('position3d')
    def _run(self):
        with base.connect() as con:
            c = con.cursor()
            #
            # 1. Create files showing the physical locations of positions with
            #    epdata, no epdata and no mutations (only positions present in
            #    the homology (i.e. no linkers)
            #
            # Get all positions with mutations (no exac)
            q = 'select distinct idx from report where pub != "exac"'
            mutations = set()
            for row in c.execute(q):
                mutations.add(row['idx'])
            # Get all positions with epdata
            epdata = set()
            for row in c.execute('select distinct idx from epdata'):
                epdata.add(row['idx'])
            # Write dimos locations of SCN5A mutations, where known
            file1 = self.data_out('dimos-locations-epdata.csv')
            file2 = self.data_out('dimos-locations-mutation.csv')
            file3 = self.data_out('dimos-locations-free.csv')
            try:
                f1 = open(file1, 'w')
                f2 = open(file2, 'w')
                f3 = open(file3, 'w')
                c1 = self.csv_writer(f1)
                c2 = self.csv_writer(f2)
                c3 = self.csv_writer(f3)
                data = ['idx','x','y','z','r','t']
                c1.writerow(data)
                c2.writerow(data)
                c3.writerow(data)
                print('Writing data to ' + file1)
                print('            and ' + file2)
                print('            and ' + file3)
                for row in c.execute('select * from scn5a_dimos_locations'):
                    # Store data
                    idx = row['idx']
                    data = [idx, row['x'], row['y'],row['z'],row['r'],row['t']]
                    if idx in epdata:
                        c1.writerow(data)
                    elif idx in mutations:
                        c2.writerow(data)
                    else:
                        c3.writerow(data)
            finally:
                f1.close()
                f2.close()
                f3.close()
            #
            # 2. Create a file with dimos positions of mutations and the
            #    associated dvi and dva
            #
            # Get midpoint shifts
            q = 'select * from epdata_filtered'
            q += ' where dva is not null and dvi is not null'
            names = {}
            dvi = {}
            dva = {}
            for row in c.execute(q):
                names[row['idx']] = row['old']+str(row['idx'])+row['new']
                dvi[row['idx']] = row['dvi']
                dva[row['idx']] = row['dva']
            # Write file with distance to pore and midpoint shifts
            file1 = self.data_out('voltage-shift-dimos-locations.csv')
            print('Writing data to ' + file1)
            with open(file1, 'w') as f1:
                c1 = self.csv_writer(f1)
                c1.writerow([
                    'name',
                    'r',
                    'dva',
                    'dvi',
                    'dva_abs',
                    'dvi_abs',
                    'sum_abs',
                    'window',
                    ])
                for row in c.execute('select * from scn5a_dimos_locations'):
                    idx = row['idx']
                    # Get epdata or skip
                    try:
                        name = names[idx]
                        da = dva[idx]
                        di = dvi[idx]
                    except KeyError:
                        continue
                    c1.writerow([
                        name,
                        row['r'],
                        da,
                        di,
                        abs(di),
                        abs(da),
                        abs(di) + abs(di),
                        di - da,
                        ])
            #
            # 3. Create a file with dimos positions, grouped by SCN5A region
            #
            # Get regions, scn5a idx to region mapping
            regions = []
            scn5a_regions = {}
            for row in c.execute('select * from region order by start'):
                regions.append(row['name'])
                for idx in range(row['start'], 1 + row['end']):
                    scn5a_regions[idx] = row['name']
            # Create files
            path = self.data_out('dimos_location_regions')
            if not os.path.isdir(path):
                os.makedirs(path)
            files = []
            csvs = {}
            try:
                # Open file per region
                for region in regions:
                    filename = region.replace(' ', '-')
                    filename = os.path.join(path, region + '.csv')
                    print('Writing data to ' + filename)
                    f = open(filename, 'w')
                    files.append(f)
                    csv = self.csv_writer(f)
                    csvs[region] = csv
                    csv.writerow(['idx', 'x', 'y', 'z'])
                # Get and write data
                q = 'select * from scn5a_dimos_locations order by idx'
                for row in c.execute(q):
                    idx = row['idx']
                    # Write to correct file
                    csv = csvs[scn5a_regions[idx]]
                    csv.writerow([idx, row['x'], row['y'], row['z']])
            finally:
                for f in files:
                    f.close()
if __name__ == '__main__':
    t = base.TaskRunner()
    t.add_tasks(tasks())
    t.run()


import cv2
import numpy as np

def detectAndDisplay(frame, mask_img, args, i):
    face_cascade_name = args.face_cascade
    eyes_cascade_name = args.eyes_cascade
    face_cascade = cv2.CascadeClassifier()
    eyes_cascade = cv2.CascadeClassifier()
    #-- 1. Load the cascades
    if not face_cascade.load(cv2.samples.findFile(face_cascade_name)):
        print('--(!)Error loading face cascade')
        exit(0)
    if not eyes_cascade.load(cv2.samples.findFile(eyes_cascade_name)):
        print('--(!)Error loading eyes cascade')
        exit(0)
    frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    frame_gray = cv2.equalizeHist(frame_gray)
    mask_h = mask_img.shape[0]
    mask_w = mask_img.shape[1]

    #-- Detect faces
    faces = face_cascade.detectMultiScale(frame_gray)
    for (x,y,w,h) in faces:
        center = (x + w//2, y + h//2)
        #translation matrix to put mask into the center of the face
        mat = np.array([[1.0, 0.0, center[0] - mask_h//4],
                        [0.0, 1.0, center[1] - mask_w//4]]).reshape(2,3)
        #frame = cv2.ellipse(frame, center, (w//2, h//2), 0, 0, 360, (255, 0, 255), 4)
        mask_img = cv2.resize(mask_img, (mask_h//2, mask_w//2))
        frame = cv2.warpAffine(mask_img, M=mat, dsize=(frame.shape[1], frame.shape[0]), borderMode=cv2.BORDER_TRANSPARENT, dst=frame,)
        faceROI = frame_gray[y:y+h,x:x+w]
        #-- In each face, detect eyes
    cv2.imshow('Capture - Face detection', frame)
    cv2.imwrite(str(i) + '.png', frame)