"""Vocabulary for theme-based transformer

    Author: Ian Shih
    Email: yjshih23@gmail.com
    Date: 2021/11/03

"""
import math

import miditoolkit
import numpy as np
from miditoolkit.midi import containers as ct
from miditoolkit.midi import parser as mid_parser

class Vocab(object):
    def __init__(self):
        """initialize some vocabulary settings"""

        # split each beat into 4 subbeats
        self.q_beat = 4

        # dictionary for matching token ID to name and the other way around.
        self.token2id = {}
        self.id2token = {}

        # midi pitch number : 1 ~ 127 (highest pitch)
        self._pitch_bins = np.arange(start=1, stop=128)

        # duration tokens 1~64 of self.q_beat
        self._duration_bins = np.arange(start=1, stop=self.q_beat * 16 + 1)

        # velocity tokens 1~127 (corressponding to midi format)
        self._velocity_bins = np.arange(start=1, stop=127)

        # tempo tokens 17~197 (determined from our dataset)
        self._tempo_bins = np.arange(start=17, stop=197, step=3)

        # position(subbeat) tokens 0~15, indicate the relative position with in a bar
        self._position_bins = np.arange(start=0, stop=16)

        self.n_tokens = 0

        self.token_type_base = {}

        self.tracks = ["MELODY", "BRIDGE", "PIANO"]

        self.build()

        # vocab
        # Note-On (129) : 0 (padding) ,1 ~ 127(highest pitch) , 128 (rest)
        # Note-Duration : 1 ~ 16 beat * 3
        # min resulution 1/12 notes

    def build(self):
        """build our vocab"""
        self.token2id = {}
        self.id2token = {}

        self.n_tokens = 0

        self.token2id["padding"] = 0
        self.n_tokens += 1

        # Note related tokens==================================================================
        # Create Note-On tokens for each track
        for track in self.tracks:
            # Note-On
            self.token_type_base = {"Note-On-{}".format(track): 1}
            for i in self._pitch_bins:
                self.token2id["Note-On-{}_{}".format(track, i)] = self.n_tokens
                self.n_tokens += 1

        # Create Note-Duration tokens for each track
        for track in self.tracks:
            # Note-Duration
            self.token_type_base["Note-Duration-{}".format(track)] = self.n_tokens
            for note_dur in self._duration_bins:
                self.token2id[
                    "Note-Duration-{}_{}".format(track, note_dur)
                ] = self.n_tokens
                self.n_tokens += 1

        # Create Note-Velocity tokens for each track
        for track in self.tracks:
            # Note-Velocity
            self.token_type_base["Note-Velocity-{}".format(track)] = self.n_tokens
            for vel in self._velocity_bins:
                self.token2id["Note-Velocity-{}_{}".format(track, vel)] = self.n_tokens
                self.n_tokens += 1

        # Metric related tokens==================================================================
        # Tempo
        self.token_type_base["Tempo"] = self.n_tokens
        for tmp in self._tempo_bins:
            self.token2id["Tempo_{}".format(tmp)] = self.n_tokens
            self.n_tokens += 1

        # Positions
        self.token_type_base["Position"] = self.n_tokens
        for pos in self._position_bins:
            self.token2id["Position_{}".format(pos)] = self.n_tokens
            self.n_tokens += 1

        # Bar
        self.token_type_base["Bar"] = self.n_tokens
        self.token2id["Bar"] = self.n_tokens
        self.n_tokens += 1

        # Theme related tokens==================================================================
        # Phrase annotation (not used in our final paper)
        self.token_type_base["Phrase"] = self.n_tokens
        self.token2id["Phrase_Start"] = self.n_tokens
        self.n_tokens += 1

        self.token2id["Phrase_End"] = self.n_tokens
        self.n_tokens += 1

        # Theme annotation
        self.token_type_base["Theme"] = self.n_tokens
        self.token2id["Theme_Start"] = self.n_tokens
        self.n_tokens += 1

        self.token2id["Theme_End"] = self.n_tokens
        self.n_tokens += 1

        for w, v in self.token2id.items():
            self.id2token[v] = w

        self.n_tokens = len(self.token2id)

    def getPitch(self, input_event):
        """Return corresponding note pitch
        if input_event is not a note, it returns -1

        Args:
            input_event (str or int): REMI Event Name or vocab ID
        """
        if isinstance(input_event, int):
            input_event = self.id2token[input_event]
        elif isinstance(input_event, str):
            pass
        else:
            try:
                input_event = int(input_event)
                input_event = self.id2token[input_event]
            except:
                raise TypeError(
                    "input_event should be int or str, input_event={}, type={}".format(
                        input_event, type(input_event)
                    )
                )

        if not input_event.startswith("Note-On"):
            return -1

        assert (
            int(input_event.split("_")[1]) >= 1
            and int(input_event.split("_")[1]) <= 127
        )
        return int(input_event.split("_")[1])

    def midi2REMI(
        self,
        midi_path,
        quantize=True,
        trim_intro=True,
        trim_outro=True,
        include_bridge=False,
        extend_theme=False,
        bar_first=False,
        theme_annotations=True,
        verbose=False,
    ):
        """convert midi file to token representation

        Args:
            midi_path (str): the path of input midi file
            trim_intro (bool, optional): ignore the intro part of the song. Defaults to True.
            trim_outro (bool, optional): ignore the outro part of the song. Defaults to True.
            include_bridge (bool, optional): ignore the intro part of the song. Defaults to False.
            extend_theme (bool, optional): extend the theme region to at least MIN_MEL_NOTES=8 notes. Defaults to False.
            bar_first (bool, optional): place Bar token in front of Theme-Start token. Defaults to False.
            theme_annotations (bool, optional): including theme-realted tokens. Defaults to True.
            verbose (bool, optional): print some message. Defaults to False.

        Returns:
            list: sequence of tokens
        """
        MIN_MEL_NOTES = 8
        midi_obj = mid_parser.MidiFile(midi_path)
        # calculate the min step (in ticks) for REMI representation
        min_step = midi_obj.ticks_per_beat * 4 / 16

        # quantize
        if quantize:
            for i in range(len(midi_obj.instruments)):
                for n in range(len(midi_obj.instruments[i].notes)):
                    midi_obj.instruments[i].notes[n].start = int(
                        int(midi_obj.instruments[i].notes[n].start / min_step)
                        * min_step
                    )
                    midi_obj.instruments[i].notes[n].end = int(
                        int(midi_obj.instruments[i].notes[n].end / min_step) * min_step
                    )

        if theme_annotations:
            # select theme info track
            theme_boundary_track = list(
                filter(lambda x: x.name == "theme info track", midi_obj.instruments)
            )
            assert len(theme_boundary_track) == 1

        # parsing notes in each tracks (ignore BRIDGE)
        notesAndtempos = []
        midi_obj.instruments[0].notes = sorted(
            midi_obj.instruments[0].notes, key=lambda x: x.start
        )
        # add notes
        melody_start = sorted(midi_obj.instruments[0].notes, key=lambda x: x.start)[
            0
        ].start
        melody_end = sorted(midi_obj.instruments[0].notes, key=lambda x: x.start)[
            -1
        ].end
        for i in range(3):
            if not include_bridge and midi_obj.instruments[i].name == "BRIDGE":
                continue
            if midi_obj.instruments[i].name == "Theme info track":
                continue

            notes = midi_obj.instruments[i].notes
            for n in notes:
                # assert (trim_intro and n.start>=melody_start or not trim_intro)
                if trim_intro and n.start >= melody_start or not trim_intro:
                    if trim_outro and n.start <= melody_end or not trim_outro:
                        notesAndtempos.append(
                            {
                                "priority": i + 1,
                                "priority_1": n.pitch,
                                "start_tick": n.start,
                                "obj_type": "Note-{}".format(
                                    midi_obj.instruments[i].name
                                ),
                                "obj": n,
                            }
                        )
        # add tempos
        for t in midi_obj.tempo_changes:
            # assert (trim_intro and t.time>=melody_start or not trim_intro)
            if trim_intro and t.time >= melody_start or not trim_intro or trim_intro:
                if trim_outro and t.time <= melody_end or not trim_outro:
                    notesAndtempos.append(
                        {
                            "priority": 0,
                            "priority_1": 0,
                            "start_tick": t.time,
                            "obj_type": "Tempo",
                            "obj": t,
                        }
                    )

        if (
            trim_intro
            and len([x for x in midi_obj.tempo_changes if x.time == melody_start]) == 0
        ):
            t = [
                x
                for x in sorted(midi_obj.tempo_changes, key=lambda z: z.time)
                if x.time < melody_start
            ]
            if not len(t) == 0:
                t = t[-1]
                notesAndtempos.append(
                    {
                        "priority": 0,
                        "priority_1": 0,
                        "start_tick": melody_start,
                        "obj_type": "Tempo",
                        "obj": t,
                    }
                )
        notesAndtempos = sorted(
            notesAndtempos,
            key=lambda x: (x["start_tick"], x["priority"], -x["priority_1"]),
        )

        if theme_annotations:
            theme_boundary_track = theme_boundary_track[0]
            theme_boundary_pitch = min([x.pitch for x in theme_boundary_track.notes])
            theme_boundaries = [
                [x.start, x.end]
                for x in theme_boundary_track.notes
                if x.pitch == theme_boundary_pitch
            ]
            assert not len(theme_boundaries) == 0
            if verbose:
                print(theme_boundaries)

        if extend_theme:
            # extend theme region 8~9
            for b_i, b in enumerate(theme_boundaries[:-1]):

                melody_notes = [
                    x
                    for x in midi_obj.instruments[0].notes
                    if x.start >= b[0] and x.start < theme_boundaries[b_i + 1][0]
                ]
                cur_bound = 0
                for x in melody_notes:
                    if x.start < b[1]:
                        cur_bound += 1
                    else:
                        break
                if cur_bound + 1 >= MIN_MEL_NOTES:
                    continue
                # try to extend
                extend_idx = min(MIN_MEL_NOTES, len(melody_notes)) - 1
                theme_boundaries[b_i][1] = melody_notes[extend_idx].end

        b_i = 0
        in_theme = False
        # group
        bar_segments = []
        bar_ticks = midi_obj.ticks_per_beat * 4
        if verbose:
            print("Bar tick length: {}".format(bar_ticks))

        for bar_start_tick in range(0, notesAndtempos[-1]["start_tick"], bar_ticks):
            if verbose:
                print(
                    "Bar {} at tick: {}".format(
                        bar_start_tick // bar_ticks, bar_start_tick
                    )
                )
            bar_end_tick = bar_start_tick + bar_ticks
            current_bar = []
            bar_objs = list(
                filter(
                    lambda x: x["start_tick"] >= bar_start_tick
                    and x["start_tick"] < bar_end_tick,
                    notesAndtempos,
                )
            )
            bar_objs.insert(0, {"start_tick": -1})

            if not bar_first:
                if (
                    theme_annotations
                    and not in_theme
                    and theme_boundaries[b_i][0] == bar_start_tick
                ):
                    current_bar.append("Theme_Start")
                    in_theme = True
                    if verbose:
                        print("Theme start")

                if (
                    not in_theme
                    and trim_intro
                    and bar_start_tick + bar_ticks < melody_start
                ):
                    if verbose:
                        print("into trimmed")
                    continue

                current_bar.append("Bar")
            else:
                if (
                    not in_theme
                    and trim_intro
                    and bar_start_tick + bar_ticks < melody_start
                ):
                    if verbose:
                        print("into trimmed")
                    continue

                current_bar.append("Bar")

                if (
                    theme_annotations
                    and not in_theme
                    and theme_boundaries[b_i][0] == bar_start_tick
                ):
                    current_bar.append("Theme_Start")
                    in_theme = True
                    if verbose:
                        print("Theme start")

            for i, obj in enumerate(bar_objs):
                if obj["start_tick"] == -1:
                    continue
                if not obj["start_tick"] == bar_objs[i - 1]["start_tick"]:
                    # insert Position Event
                    pos = (
                        (obj["start_tick"] - bar_start_tick)
                        / midi_obj.ticks_per_beat
                        * self.q_beat
                    )
                    pos_index = np.argmin(
                        abs(pos - self._position_bins)
                    )  # use the closest position
                    pos = self._position_bins[pos_index]
                    current_bar.append("Position_{}".format(pos))

                if obj["obj_type"].startswith("Note"):
                    track_name = obj["obj_type"].split("-")[1].upper()
                    # add pitch
                    current_bar.append(
                        "Note-On-{}_{}".format(track_name, obj["obj"].pitch)
                    )
                    # add duration
                    dur = (
                        (obj["obj"].end - obj["obj"].start)
                        / midi_obj.ticks_per_beat
                        * self.q_beat
                    )
                    dur_index = np.argmin(
                        abs(dur - self._duration_bins)
                    )  # use the closest position
                    dur = self._duration_bins[dur_index]
                    current_bar.append("Note-Duration-{}_{}".format(track_name, dur))
                    # add velocity
                    vel_index = np.argmin(
                        abs(obj["obj"].velocity - self._velocity_bins)
                    )  # use the closest position
                    vel = self._velocity_bins[vel_index]
                    current_bar.append("Note-Velocity-{}_{}".format(track_name, vel))
                elif obj["obj_type"].startswith("Tempo"):
                    # tempo
                    tmp_index = np.argmin(
                        abs(obj["obj"].tempo - self._tempo_bins)
                    )  # use the closest position
                    tmp = self._tempo_bins[tmp_index]
                    current_bar.append(obj["obj_type"] + "_{}".format(tmp))
                else:
                    # theme start end
                    current_bar.append(obj["obj_type"])

            if (
                theme_annotations
                and in_theme
                and theme_boundaries[b_i][1] == bar_start_tick + bar_ticks
            ):
                current_bar.append("Theme_End")
                in_theme = False
                if verbose:
                    print("Theme End")
                if not b_i == len(theme_boundaries) - 1:
                    b_i += 1
            bar_segments.extend(current_bar)

        output_ids = [self.token2id[x] for x in bar_segments]

        return output_ids

    def preprocessREMI(
        self,
        remi_sequence,
        always_include=False,
        max_seq_len=512,
        strict=True,
        verbose=False,
    ):
        """Preprocess token sequence

        slicing the sequence for training our models

        Args:
            remi_sequence (List): the music token seqeunce
            always_include (Bool): selected the data including either Theme-Start or Theme-End
            max_seq_len (Int): maximum sequence length for each data
            strict (Bool): the returning sequence should always include Theme-Start

        Return:
            {
                "src" : <corressponding theme condition>,
                "src_theme_binary_msk" : <corressponding theme condition's theme msk>,
                "tgt_segments" : <list of target sequences>,
                "tgt_segments_theme_binary_msk" : <list of target sequences theme msk>
            }
        """
        theme_binary_msk = []
        in_theme = False
        src = []
        src_theme_binary_msk = []
        for r in remi_sequence:
            if self.id2token[r] == "Theme_Start":
                in_theme = True
            elif self.id2token[r] == "Theme_End":
                in_theme = False
            theme_binary_msk.append(int(in_theme))

        for i in range(1, len(theme_binary_msk)):
            theme_binary_msk[i] = (
                theme_binary_msk[i - 1] * theme_binary_msk[i] + theme_binary_msk[i]
            )

        start_first_theme = remi_sequence.index(self.token2id["Theme_Start"])
        end_first_theme = remi_sequence.index(self.token2id["Theme_End"])
        src = remi_sequence[start_first_theme : end_first_theme + 1]
        src_theme_binary_msk = theme_binary_msk[start_first_theme : end_first_theme + 1]

        tgt_segments = []
        tgt_segments_theme_msk = []

        s = 0
        if strict:
            theme_start_pos = [
                i
                for i in range(len(remi_sequence))
                if remi_sequence[i] == self.token2id["Theme_Start"]
            ]
            for t in theme_start_pos:
                tgt_segments.append(remi_sequence[t : t + max_seq_len + 1])
                tgt_segments_theme_msk.append(theme_binary_msk[t : t + max_seq_len + 1])
        else:
            total_s = math.ceil(len(remi_sequence) / max_seq_len)
            for x in range(0, len(remi_sequence), max_seq_len):
                if always_include:
                    # if self.token2id["Theme_Start"] in remi_sequence[x:x+max_seq_len+1]:
                    if (
                        self.token2id["Theme_Start"]
                        in remi_sequence[x : x + max_seq_len + 1]
                        or self.token2id["Theme_End"]
                        in remi_sequence[x : x + max_seq_len + 1]
                    ):
                        s += 1
                        tgt_segments.append(remi_sequence[x : x + max_seq_len + 1])
                        tgt_segments_theme_msk.append(
                            theme_binary_msk[x : x + max_seq_len + 1]
                        )
                else:
                    tgt_segments.append(remi_sequence[x : x + max_seq_len + 1])
                    tgt_segments_theme_msk.append(
                        theme_binary_msk[x : x + max_seq_len + 1]
                    )

            if verbose and always_include:
                print("Include Theme Start {}/{}".format(s, total_s))

        return {
            "src": src,
            "src_theme_binary_msk": src_theme_binary_msk,
            "tgt_segments": tgt_segments,
            "tgt_segments_theme_binary_msk": tgt_segments_theme_msk,
        }

    def REMIID2midi(self, event_ids, midi_path, verbose=False):
        """convert tokens to midi file
        The output midi file will contains 3 tracks:
            MELODY : melodt notes
            PIANO : accompaniment notes
            Theme info track : notes indicating theme region (using note start tick and end tick as the boundary of theme region)
        Args:
            event_ids (list): sequence of tokens
            midi_path (str): the output midi file path
            verbose (bool, optional): print some message. Defaults to False.
        """

        # create midi file
        new_mido_obj = mid_parser.MidiFile()
        new_mido_obj.ticks_per_beat = 120

        # create tracks
        music_tracks = {}
        music_tracks["MELODY"] = ct.Instrument(program=0, is_drum=False, name="MELODY")
        music_tracks["PIANO"] = ct.Instrument(program=0, is_drum=False, name="PIANO")
        music_tracks["Theme info track"] = ct.Instrument(
            program=0, is_drum=False, name="Theme info track"
        )

        # all our generated music are 4/4
        new_mido_obj.time_signature_changes.append(miditoolkit.TimeSignature(4, 4, 0))

        ticks_per_step = new_mido_obj.ticks_per_beat / self.q_beat

        # convert tokens from id to string
        events = []
        for x in event_ids:
            events.append(self.id2token[x])

        # parsing tokens
        last_tick = 0
        current_bar_anchor = 0
        current_theme_boundary = []
        motif_label_segs = []
        idx = 0
        first_bar = True
        while idx < len(events):
            if events[idx] == "Bar":
                if first_bar:
                    current_bar_anchor = 0
                    first_bar = False
                else:
                    current_bar_anchor += new_mido_obj.ticks_per_beat * 4
                idx += 1
            elif events[idx].startswith("Position"):
                pos = int(events[idx].split("_")[1])
                last_tick = pos * ticks_per_step + current_bar_anchor
                idx += 1
            elif events[idx].startswith("Tempo"):
                tmp = pos = int(events[idx].split("_")[1])
                new_mido_obj.tempo_changes.append(
                    ct.TempoChange(tempo=int(tmp), time=int(last_tick))
                )
                idx += 1
            elif events[idx].startswith("Note"):
                track_name = events[idx].split("_")[0].split("-")[2]
                # assert track_name in music_tracks
                assert events[idx].startswith("Note-On")
                assert events[idx + 1].startswith("Note-Duration")
                assert events[idx + 2].startswith("Note-Velocity")

                if track_name == "BRIDGE":
                    track_name = "PIANO"

                new_note = miditoolkit.Note(
                    velocity=int(events[idx + 2].split("_")[1]),
                    pitch=int(events[idx].split("_")[1]),
                    start=int(last_tick),
                    end=int(int(events[idx + 1].split("_")[1]) * ticks_per_step)
                    + int(last_tick),
                )
                music_tracks[track_name].notes.append(new_note)
                idx += 3
            elif events[idx] == "Theme_Start":
                assert len(current_theme_boundary) == 0
                current_theme_boundary.append(last_tick)
                idx += 1
            elif events[idx] == "Theme_End":
                assert len(current_theme_boundary) == 1
                current_theme_boundary.append(last_tick)
                motif_label_segs.append(current_theme_boundary)
                music_tracks["Theme info track"].notes.append(
                    miditoolkit.Note(
                        velocity=1,
                        pitch=1,
                        start=int(current_theme_boundary[0]),
                        end=int(current_theme_boundary[1]),
                    )
                )
                current_theme_boundary = []
                idx += 1
            else:
                idx += 1

        # add tracks to midi file
        new_mido_obj.instruments.extend([music_tracks[ins] for ins in music_tracks])

        if verbose:
            print("Saving midi to ({})".format(midi_path))

        # save to disk
        new_mido_obj.dump(midi_path)

    def __str__(self):
        """return all tokens

        Returns:
            str: string of all tokens
        """
        ret = ""
        for w, i in self.token2id.items():
            ret = ret + "{} : {}\n".format(w, i)

        for i, w in self.id2token.items():
            ret = ret + "{} : {}\n".format(i, w)

        ret += "\nTotal events #{}".format(len(self.id2token))

        return ret

    def __repr__(self):
        """return string all token

        Returns:
            str: string of sll tokens
        """
        return self.__str__()

if __name__ == "__main__":
    # print all tokens
    myvocab = Vocab()

    print(myvocab)


from collections import deque

import numpy as np
import cv2

import chainer
from chainer import links as L

import chainerrl
from chainerrl import agents
from chainerrl.action_value import DiscreteActionValue
from chainerrl import explorers
from chainerrl import links
from chainerrl import replay_buffer

def infer(agent, state):
    gray_state = [cv2.cvtColor(s, cv2.COLOR_RGB2GRAY) for s in state]
    dqn_state = [cv2.resize(s, (84, 84), interpolation=cv2.INTER_AREA) \
                 for s in gray_state]
    input_tensor = np.array(dqn_state).astype(np.float32)
    return agent.act(input_tensor)

class Agent(object):
    def __init__(self,
                 modelpath,
                 n_actions=4,
                 n_stack_frames=4):
        # Predefined parameters.
        replay_start_size = 5 * 10 ** 4

        # Load the model.
        q_func = links.Sequence(
            links.NatureDQNHead(),
            L.Linear(512, n_actions),
            DiscreteActionValue)
        opt = chainer.optimizers.RMSpropGraves(
            lr=2.5e-4, alpha=0.95, momentum=0.0, eps=1e-2)
        opt.setup(q_func)
        rbuf = replay_buffer.ReplayBuffer(10 ** 6)
        explorer = explorers.LinearDecayEpsilonGreedy(
            start_epsilon=1.0, end_epsilon=0.1,
            decay_steps=10 ** 6,
            random_action_func=lambda: np.random.randint(n_actions))
        def phi(x):
            # Feature extractor
            return np.asarray(x, dtype=np.float32) / 255

        Agent = agents.DQN
        self._agent = Agent(q_func, opt, rbuf, gpu=-1, gamma=0.99,
                      explorer=explorer, replay_start_size=replay_start_size,
                      target_update_interval=10 ** 4,
                      clip_delta=True,
                      update_interval=4,
                      batch_accumulator='sum',
                      phi=phi)
        self._agent.load(modelpath)

        self._state = deque(
            [], maxlen=n_stack_frames)
        self._action = 0

    def get_action(self):
        return self._action

    def put_state(self, state):
        # Note: should devide this code to 2 parts:
        # putting state part and do inference part...
        self._state.append(state)
        if len(self._state) < self._state.maxlen:
            # Need to wait.
            return
        state = list(self._state)
        self._action = infer(self._agent, state)


"""
This file contains a function to decode a single synthetic tree.
"""
import pandas as pd
import numpy as np
from dgllife.model import load_pretrained
from syn_net.utils.data_utils import ReactionSet
from syn_net.utils.predict_utils import synthetic_tree_decoder, tanimoto_similarity, load_modules_from_checkpoint

# define some constants (here, for the Hartenfeller-Button test set)
nbits        = 4096
out_dim      = 256
rxn_template = 'hb'
featurize    = 'fp'
param_dir    = 'hb_fp_2_4096_256'
ncpu         = 16

# define model to use for molecular embedding
model_type   = 'gin_supervised_contextpred'
device       = 'cpu'
mol_embedder = load_pretrained(model_type).to(device)
mol_embedder.eval()

# load the purchasable building block embeddings
bb_emb = np.load('/pool001/whgao/data/synth_net/st_hb/enamine_us_emb_fp_256.npy')

# define path to the reaction templates and purchasable building blocks
path_to_reaction_file   = f'/pool001/whgao/data/synth_net/st_{rxn_template}/reactions_{rxn_template}.json.gz'
path_to_building_blocks = f'/pool001/whgao/data/synth_net/st_{rxn_template}/enamine_us_matched.csv.gz'

# define paths to pretrained modules
param_path  = f'/home/whgao/synth_net/synth_net/params/{param_dir}/'
path_to_act = f'{param_path}act.ckpt'
path_to_rt1 = f'{param_path}rt1.ckpt'
path_to_rxn = f'{param_path}rxn.ckpt'
path_to_rt2 = f'{param_path}rt2.ckpt'

# load the purchasable building block SMILES to a dictionary
building_blocks = pd.read_csv(path_to_building_blocks, compression='gzip')['SMILES'].tolist()
bb_dict         = {building_blocks[i]: i for i in range(len(building_blocks))}

# load the reaction templates as a ReactionSet object
rxn_set = ReactionSet()
rxn_set.load(path_to_reaction_file)
rxns    = rxn_set.rxns

# load the pre-trained modules
act_net, rt1_net, rxn_net, rt2_net = load_modules_from_checkpoint(
    path_to_act=path_to_act,
    path_to_rt1=path_to_rt1,
    path_to_rxn=path_to_rxn,
    path_to_rt2=path_to_rt2,
    featurize=featurize,
    rxn_template=rxn_template,
    out_dim=out_dim,
    nbits=nbits,
    ncpu=ncpu,
)

def func(emb):
    """
    Generates the synthetic tree for the input molecular embedding.

    Args:
        emb (np.ndarray): Molecular embedding to decode.

    Returns:
        str: SMILES for the final chemical node in the tree.
        SyntheticTree: The generated synthetic tree.
    """
    emb = emb.reshape((1, -1))
    try:
        tree, action = synthetic_tree_decoder(z_target=emb,
                                              building_blocks=building_blocks,
                                              bb_dict=bb_dict,
                                              reaction_templates=rxns,
                                              mol_embedder=mol_embedder,
                                              action_net=act_net,
                                              reactant1_net=rt1_net,
                                              rxn_net=rxn_net,
                                              reactant2_net=rt2_net,
                                              bb_emb=bb_emb,
                                              rxn_template=rxn_template,
                                              n_bits=nbits,
                                              max_step=15)
    except Exception as e:
        print(e)
        action = -1
    if action != 3:
        return None, None
    else:
        scores = np.array(
            tanimoto_similarity(emb, [node.smiles for node in tree.chemicals])
        )
        max_score_idx = np.where(scores == np.max(scores))[0][0]
        return tree.chemicals[max_score_idx].smiles, tree


# -*- mode: python; coding: utf-8 -*-
# Copyright 2018 Peter Williams and collaborators.
# Licensed under the MIT License.

"""This module provides helpers for doing radiative transfer integrations with
`grtrans <https://github.com/jadexter/grtrans>`_, including a simple
framework for running end-to-end tests.

In order to use this functionality, the Python module ``radtrans_integrate``
must be importable. Sadly `grtrans <https://github.com/jadexter/grtrans>`_
doesn't install itself like a regular Python package, so getting this working
can be a pain. Documenting the installation procedure is beyond the scope of
this project.

"""
from __future__ import absolute_import, division, print_function

import numpy as np
import pandas as pd

# grtrans' integration methods:
METHOD_LSODA_YES_LINEAR_STOKES = 0 # LSODA with IS_LINEAR_STOKES=1
METHOD_DELO = 1 # DELO method from Rees+ (1989ApJ...339.1093R)
METHOD_FORMAL = 2 # "formal" method = "matricant (O-matrix) method from Landi Degl'Innocenti"
METHOD_LSODA_NO_LINEAR_STOKES = 3 # LSODA with IS_LINEAR_STOKES=0 -- this is "under development" spherical stokes

def integrate_ray_formal(x, j, K):
    """Use `grtrans <https://github.com/jadexter/grtrans>`_ to integrate one ray
    using its "formal" (matricant / O-matrix) method.

    **Call signature**

    *x*
      1D array, shape (n,). Path length along ray, starting from zero, in cm.
    *j*
      Array, shape (n, 4). Emission coefficients: ``j_{IQUV}``, in that order.
    *K*
      Array, shape (n, 7). Absorption coefficients and Faraday mixing coefficients:
      ``alpha_{IQUV}, rho_{QUV}``.
    Return value
      Array of shape (4, m): Stokes intensities ``IQUV`` along parts of the
      ray with non-zero total emissivities; m <= n.

    """
    # A correct version of the fully-specified absorption matrix is in Leung+
    # 2011 (10.1088/0004-637X/737/1/21). From looking at the function
    # `radtrans_jac_form` of grtrans' `radtrans_integrate.f90` file, one can
    # see that the K vector is indeed packed in the way described above.
    from radtrans_integrate import radtrans_integrate

    n = x.size

    # The formal method doesn't do the same kind of clipping as LSODA *inside*
    # grtrans, but here we do the same clipping for consistency, and since it
    # is genuinely true that the clipped samples are superfluous.

    if np.all(j[:,0] == 0.):
        return np.zeros((4, n))

    i0 = 0
    i1 = n - 1

    while j[i0,0] == 0.:
        i0 += 1
    while j[i1,0] == 0.:
        i1 -= 1

    n = i1 + 1 - i0
    x = x[i0:i1+1]
    j = j[i0:i1+1]
    K = K[i0:i1+1]

    # OK we can go.

    radtrans_integrate.init_radtrans_integrate_data(
        METHOD_FORMAL, # method selector
        4, # number of equations
        n, # number of input data points
        n, # number of output data points
        10., # maximum optical depth; not used by "formal"
        1., # maximum absolute step size; not used by "formal"
        0.1, # absolute tolerance; not used by "formal"
        0.1, # relative tolerance; not used by "formal"
        1e-2, # "thin" parameter for DELO method; not used by "formal"
        1, # maximum number of steps; not used by "formal"
    )

    try:
        tau = x # not used by "formal"
        radtrans_integrate.integrate(x[::-1], j[::-1], K[::-1], tau[::-1], 4)
        i = radtrans_integrate.intensity.copy()
    finally:
        # If we exit without calling this, the next init call causes an abort
        radtrans_integrate.del_radtrans_integrate_data()
    return i

def integrate_ray_lsoda(x, j, K, atol=1e-8, rtol=1e-6, max_step_size=None,
                        frac_max_step_size=1e-3, max_steps=100000):
    """Use `grtrans <https://github.com/jadexter/grtrans>`_ to integrate one ray
    using its LSODA method.

    **Call signature**

    *x*
      1D array, shape (n,). Path length along ray, starting from zero, in cm.
    *j*
      Array, shape (n, 4). Emission coefficients: ``j_{IQUV}``, in that order.
    *K*
      Array, shape (n, 7). Absorption coefficients and Faraday mixing coefficients:
      ``alpha_{IQUV}, rho_{QUV}``.
    *atol*
      Some kind of tolerance parameter.
    *rtol*
      Some kind of tolerance parameter.
    *max_step_size*
      The maximum absolute step size. Overrides *frac_max_step_size*.
    *frac_max_step_size*
      If *max_step_size* is not specified, the maximum step size passed to the
      integrator is ``x.max()`` multiplied by this parameter. Experience shows
      that (for LSODA at least) this parameter must be pretty small to get
      good convergence!
    *max_steps*
      The maximum number of steps to take.
    Return value
      Array of shape (4, m): Stokes intensities IQUV along parts of the ray with
      non-zero total emissivities; m <= n.

    """
    n = x.size

    if max_step_size is None:
        max_step_size = frac_max_step_size * x.max()

    # the LSODA method clips its input arrays based on "tau" and zero emission
    # coefficients. It's hard for us to find out how it clipped, though, so we
    # reproduce its logic. LSODA doesn't use "tau" for anything besides this
    # clipping, so we pass it all zeros.

    if np.all(j[:,0] == 0.):
        return np.zeros((4, n))

    i0 = 0
    i1 = n - 1

    while j[i0,0] == 0.:
        i0 += 1
    while j[i1,0] == 0.:
        i1 -= 1

    n = i1 + 1 - i0
    x = x[i0:i1+1]
    j = j[i0:i1+1]
    K = K[i0:i1+1]

    # OK we can go.

    radtrans_integrate.init_radtrans_integrate_data(
        METHOD_LSODA_YES_LINEAR_STOKES, # method selector
        4, # number of equations
        n, # number of input data points
        n, # number of output data points
        10., # maximum optical depth; defused here (see comment above)
        max_step_size, # maximum absolute step size
        atol, # absolute tolerance
        rtol, # relative tolerance
        1e-2, # "thin" parameter for DELO method ... to be researched
        max_steps, # maximum number of steps
    )

    try:
        tau = np.zeros(n)
        radtrans_integrate.integrate(x[::-1], j, K, tau, 4)
        i = radtrans_integrate.intensity.copy()
    finally:
        # If we exit without calling this, the next init call causes an abort
        radtrans_integrate.del_radtrans_integrate_data()
    return i

def integrate(d, coeffs, psi):
    """Integrate a ray with `grtrans <https://github.com/jadexter/grtrans>`_,
    using reasonable defaults.

    **Call signature**

    *d*
      An array giving the location of each sample along the ray, starting from zero, in cm.
    *coeffs*
      An array of shape (N, 8) of RT coefficients in the basis where the
      Stokes U coefficients are always zero. Such arrays are returned by
      :meth:`neurosynchro.impl.PhysicalApproximator.compute_all_nontrivial`.
    *psi*
      An array of angles between the local magnetic field and the observer’s Stokes U
      axis, in radians.
    Return value
      An array of shape (4,), giving the Stokes IQUV at the end of the ray.

    This function is mainly intended to test what happens if the passed-in
    coefficients are slightly different due to the neural network
    approximation. So we don't provide many knobs or diagnostics here.

    """
    from . import detrivialize_stokes_basis

    xformed = detrivialize_stokes_basis(coeffs, psi)
    j = np.empty(coeffs.shape[:-1] + (4,))
    j[...,0] = xformed[...,0]
    j[...,1] = xformed[...,2]
    j[...,2] = xformed[...,4]
    j[...,3] = xformed[...,6]
    K = np.empty(coeffs.shape[:-1] + (7,))
    K[...,0] = xformed[...,1]
    K[...,1] = xformed[...,3]
    K[...,2] = xformed[...,5]
    K[...,3:] = xformed[...,7:]

    iquv = integrate_ray_formal(d, j, K)
    return iquv[:,-1]

def make_parser(ap=None):
    if ap is None:
        ap = argparse.ArgumentParser()

    ap.add_argument('--frequency', '-f', type=float, metavar='<ghz>',
                    default=10.0, help='The observing frequency to simulate')
    ap.add_argument('nndir', type=str, metavar='<nndir>',
                    help='The path to the neural-net directory.')
    ap.add_argument('testdata', type=str, metavar='<testdata>',
                    help='A file containing test data')
    return ap

def grtrans_cli(settings):
    from pwkit import cgs
    from pwkit.cli import die
    from time import time
    from .impl import PhysicalApproximator, hardcoded_nu_ref, hardcoded_ne_ref

    # Read and validate the test dataset.

    testdata = pd.read_table(settings.testdata)

    psi = testdata.get('psi(meta)')
    if psi is None:
        die('the test dataset must contain a column of field-to-Stokes-U angles \"psi(meta)\"')

    d = testdata.get('d(meta)')
    if d is None:
        die('the test dataset must contain a column of integration path lengths \"d(meta)\"')

    n_e = testdata.get('n_e(meta)')
    if n_e is None:
        die('the test dataset must contain a column of particle densities \"n_e(meta)\"')

    time_ms = testdata.get('time_ms(meta)')
    if time_ms is None:
        die('the test dataset must contain a column of computation times \"time_ms(meta)\"')

    s = None
    theta = None
    others = {}

    for col in testdata.columns:
        if col.startswith('s('):
            s = testdata[col]
        elif col.startswith('theta('):
            theta = testdata[col]
        elif col.endswith('(lin)') or col.endswith('(log)'):
            others[col.split('(')[0]] = testdata[col]

    if s is None:
        die('the test dataset must have an input parameter of the harmonic number \"s\"')

    if theta is None:
        die('the test dataset must have an input parameter of the field-to-LOS angle \"theta\"')

    # Get the coefficients into physical units, packed in our standard format.

    nu_hz = settings.frequency * 1e9
    freq_scale = nu_hz / hardcoded_nu_ref
    n_e_scale = n_e / hardcoded_ne_ref

    coeffs = np.empty((psi.size, 8))
    coeffs[...,0] = testdata['j_I(res)'] * freq_scale
    coeffs[...,1] = testdata['alpha_I(res)'] / freq_scale
    coeffs[...,2] = testdata['j_Q(res)'] * freq_scale
    coeffs[...,3] = testdata['alpha_Q(res)'] / freq_scale
    coeffs[...,4] = testdata['j_V(res)'] * freq_scale
    coeffs[...,5] = testdata['alpha_V(res)'] / freq_scale
    coeffs[...,6] = testdata['rho_Q(res)'] / freq_scale
    coeffs[...,7] = testdata['rho_V(res)'] / freq_scale
    coeffs *= n_e_scale.values.reshape((-1, 1))

    # Ground truth:

    iquv_precise = integrate(d, coeffs, psi)
    ctime_precise = time_ms.sum()
    print('Precise computation: I={:.4e}  Q={:.4e}  U={:.4e}  V={:.4e}  calc_time={:.0f} ms'.format(
          iquv_precise[0], iquv_precise[1], iquv_precise[2], iquv_precise[3], ctime_precise
    ))

    # Now set up the approximator and do the same thing. (Note that often the
    # timing seems backwards, because the time spent doing the precise
    # calculation has already been spent, whereas we have a lot of overhead to
    # set up the neural networks.)

    B = 2 * np.pi * cgs.me * cgs.c * nu_hz / (s * cgs.e)
    approx = PhysicalApproximator(settings.nndir)
    t0 = time()
    coeffs, oos = approx.compute_all_nontrivial(nu_hz, B, n_e, theta, **others)
    ctime_approx = 1000 * (time() - t0)

    if np.any(oos != 0):
        print('WARNING: some of the approximations were out-of-sample')

    iquv_approx = integrate(d, coeffs, psi)
    print('Approx. computation: I={:.4e}  Q={:.4e}  U={:.4e}  V={:.4e}  calc_time={:.0f} ms'.format(
          iquv_approx[0], iquv_approx[1], iquv_approx[2], iquv_approx[3], ctime_approx
    ))

    acc = np.abs((iquv_approx - iquv_precise) / iquv_precise)
    print('Accuracy: I={:.3f}  Q={:.3f}  U={:.3f}  V={:.3f}'.format(acc[0], acc[1], acc[2], acc[3]))

    print('Speedup: {:.1f}'.format(ctime_precise / ctime_approx))


#!/usr/bin/env python
from __future__ import division
from __future__ import print_function
from builtins import zip
from builtins import str
import sys, os, re
import argparse
from utils.labeler import read_attack_file, ip2int, int2ip, checkIPsEqual
from utils.time_functions import datetime_to_tstamp
import numpy as np
import matplotlib
import matplotlib.pyplot as plt

def check_results(results_file, attacks_file, threshold, make_plots, table_thresh):
    threshold_vals = parse_threshold(threshold)
    attack_list, num_unique_attacks = read_attack_file(attacks_file)
    raw_results = read_results(results_file)
    final_results = get_final_results(raw_results, attack_list)
    pc_attacks_detected = []
    num_TP_per_day = []
    num_FP = []
    num_FP_per_day = []
    f1s = []
    for thresh in threshold_vals:
        total_unique_TP, total_TP_per_day, total_FP, total_FP_per_day = classify_results(
            final_results, thresh
        )
        pc_attacks_detected.append(
            float(total_unique_TP) / float(num_unique_attacks) * 100
        )
        num_FP.append(total_FP)
        num_FP_per_day.append(total_FP_per_day)
        num_TP_per_day.append(total_TP_per_day)
        recall = total_unique_TP / num_unique_attacks
        total_pos = total_unique_TP + total_FP
        precision = total_unique_TP / total_pos if total_pos > 0 else 0
        if (recall + precision) > 0:
            f1s.append(2 * (recall * precision) / (recall + precision))
        else:
            f1s.append(0.0)

    data = {}
    data["threshold_vals"] = threshold_vals
    data["pc_attacks_detected"] = pc_attacks_detected
    data["num_FP"] = num_FP
    data["num_FP_per_day"] = num_FP_per_day
    data["num_TP_per_day"] = num_TP_per_day
    data["f1s"] = f1s

    if make_plots:
        plot_results(data)
    print_results(final_results, data, table_thresh)

    return data

def plot_results(data):
    # plot the pc_attacks_detected vs threshold_vals
    plt.figure()
    plt.plot(data["threshold_vals"], data["pc_attacks_detected"], marker="o")
    plt.xlabel("Threshold Values")
    plt.ylabel("% of Total Attacks Detected")
    #    plt.title('% Total Attacks Detected vs. Threshold Value')

    # plot the total # false positives vs threshold values
    plt.figure()
    plt.semilogy(data["threshold_vals"], data["num_FP"], marker="o")
    plt.xlabel("Threshold Values")
    plt.ylabel("Total # False Positives")
    #    plt.title('Total # False Positives vs. Threshold Value')

    plt.figure()
    plt.semilogx(data["num_FP"], data["pc_attacks_detected"], marker="o", linewidth=5)
    plt.xlabel("Total # False Positives")
    plt.ylabel("% of Total Attacks Detected")
    plt.yticks([0, 10, 20, 30, 40, 50])
    plt.xlim(max(data["num_FP"]), 0)
    plt.subplots_adjust(bottom=0.2)

    font = {"family": "normal", "weight": "bold", "size": 35}
    matplotlib.rc("font", **font)

    plt.show()

def parse_threshold(threshold):
    """
    Inputs:
      - threshold string of the form: start:stop:num_points
    Returns:
      - numpy array of threshold values
    """
    threshold_fmat = r"(?P<start>[\d\.]*):(?P<stop>[\d\.]*):(?P<num_points>\d*)$"
    match = re.search(threshold_fmat, threshold)
    if match is None:
        print("ERROR: invalid threshold specified", file=sys.stderr)
        sys.exit(1)
    try:
        start = float(match.group("start"))
        stop = float(match.group("stop"))
        num_points = int(match.group("num_points"))
    except ValueError as e:
        print("ERROR: invalid threshold specified", file=sys.stderr)
        sys.exit(1)
    return np.linspace(start, stop, num_points)

def get_final_results(raw_results, attack_list):
    """
    Inputs:
      - raw_results : list of dictionaries from reading results.csv file
      - attack_list : list of dictionaries of all attacks from attack file

    Outputs:
      - final_results : raw_results appended with attack info and pruned so that
          results that detect duplicate attacks are removed
    """
    detected_attackIDs = []
    final_results = []
    # sort the raw_results by score from high to low so that only the
    # highest scoring of any duplicate attacks will appear in the final_results
    raw_results = sorted(raw_results, key=getScoreVal)
    for result in raw_results:
        isAttack, attackID, attackName = get_attack_info(result, attack_list)
        result["attackID"] = attackID
        result["isAttack"] = isAttack
        result["attack_name"] = attackName
        if isAttack and attackID not in detected_attackIDs:
            # record this attack ID so that we know we've seen it before
            detected_attackIDs.append(attackID)
            # don't include any duplicate attack detections in the final results
            final_results.append(result)
        elif not isAttack:
            # include all results that do not detect an attack in the final results
            final_results.append(result)
    return final_results

def classify_results(final_results, threshold):
    """
    Inputs:
      - final_results : list of dictionaries of results, that contain info as to whether or not
          this is an actual attack

    Returns:
      - total_unique_TP : total number unique attacks classified as attacks
      - total_FP : total number of false positives
    """
    total_unique_TP = 0
    total_TP_per_day = {}
    total_FP = 0
    total_FP_per_day = {}
    for result in final_results:
        if result["score"] > threshold:
            if result["isAttack"]:
                total_unique_TP += 1
                if result["date"] not in list(total_TP_per_day.keys()):
                    total_TP_per_day[result["date"]] = 1
                else:
                    total_TP_per_day[result["date"]] += 1
            else:
                total_FP += 1
                if result["date"] not in list(total_FP_per_day.keys()):
                    total_FP_per_day[result["date"]] = 1
                else:
                    total_FP_per_day[result["date"]] += 1
    return total_unique_TP, total_TP_per_day, total_FP, total_FP_per_day

def get_attack_info(result_dic, attack_list, leeway=60):
    """
    Inputs:
      - result_dic : dictionary of one result
      - attack_list : list of attacks

    Returns:
      - isAttack - True if the result corresponds to an actual attack
      - attackID - string of attack ID, '' if not attack
      - attackName - string of attack name, '' if not an attack
    """
    isAttack = False
    attackID = ""
    attackName = ""
    tstamp = result_dic["timestamp"]
    dstIP = result_dic["dstIP"]
    for attack in attack_list:
        if (
            tstamp >= (attack["range"][0] - leeway)
            and tstamp <= (attack["range"][1] + leeway)
            and checkIPsEqual(dstIP, attack["dstIP"])
        ):
            isAttack = True
            attackID = attack["ID"]
            attackName = attack["name"]
            return isAttack, attackID, attackName
    return isAttack, attackID, attackName

def getScoreVal(item):
    return -item["score"]

#  def print_results(final_results, data, pthresh=1e-8):
def print_results(final_results, data, pthresh):
    threshold_vals = data["threshold_vals"]
    pc_attacks_detected = data["pc_attacks_detected"]
    num_FP = data["num_FP"]
    num_FP_per_day = data["num_FP_per_day"]
    num_TP_per_day = data["num_TP_per_day"]
    f1s = data["f1s"]
    detected_attackNames = []
    title_fmat = "{: >12} {: >10} {: >18}  {: >6} {: >6} {: >25} {: >25}"
    line_fmat = "{: >12} {: >10} {: >18}  {:.6f} {: >4} {: >25} {: >25}"
    #  title_fmat = '{: >12} {: >10} {: >18}  {: >6} {: >6} {: >25}'
    #  line_fmat = '{: >12} {: >10} {: >18}  {:.6f} {: >4} {: >25}'
    count = 0
    print("For threshold = {}".format(pthresh))
    print("Top 20 scores: ")
    print(
        title_fmat.format(
            "Date",
            "Time",
            "Dest IP Addr",
            "Score",
            "Det",
            "Attack Name",
            "Most Anom Field",
        )
    )
    #  print title_fmat.format('Date', 'Time', 'Dest IP Addr', 'Score', 'Det', 'Attack Name')
    for res in final_results:
        if res["isAttack"] and res["score"] >= pthresh:
            class_type = "TP"
        elif not res["isAttack"] and res["score"] >= pthresh:
            class_type = "FP"
        else:
            class_type = "N"
        if count < 20:
            print(
                line_fmat.format(
                    res["date"],
                    res["time"],
                    res["dstIP"],
                    res["score"],
                    class_type,
                    res["attack_name"],
                    res["anom_field"]
                    + " "
                    + str(int(float(res["anom_field_pc"]) * 100))
                    + "%",
                )
            )
            #  print line_fmat.format(res['date'], res['time'], res['dstIP'], res['score'], class_type, res['attack_name'])
            count += 1
        if res["attack_name"] not in detected_attackNames and res["score"] >= pthresh:
            detected_attackNames.append(res["attack_name"])
    print("# detected attackNames = {}".format(len(detected_attackNames)))
    print("Detection Results:")
    print("------------------")
    line_fmat = "{: >25} {: >25} {: >25} {: >25}"
    print(
        line_fmat.format(
            "threshold val", "% attacks detected", "Total Num FP", "F1 Score"
        )
    )
    for thresh, pc_detect, fp, f1 in zip(
        threshold_vals, pc_attacks_detected, num_FP, f1s
    ):
        print(line_fmat.format(thresh, pc_detect, fp, f1))

    for thresh, fp_per_day_dic, tp_per_day_dic in zip(
        threshold_vals, num_FP_per_day, num_TP_per_day
    ):
        print("threshold = {}, False Positives = {}".format(thresh, fp_per_day_dic))
        print("threshold = {}, True Positives = {}".format(thresh, tp_per_day_dic))

def read_results(results_file):
    result_fmat = r"(?P<date>.*),(?P<time>.*),(?P<dstIP>.*),(?P<score>.*),(?P<anom_field>.*),(?P<anom_field_pc>[\d\.]*)[\r\n$]"
    #  result_fmat = r'(?P<date>.*),(?P<time>.*),(?P<dstIP>.*),(?P<score>.*)[\r\n$]'
    with open(results_file) as f:
        contents = f.read()
    matches = re.finditer(result_fmat, contents, re.M)
    results = [m.groupdict() for m in matches]
    for dic in results:
        dic["timestamp"] = datetime_to_tstamp(dic["date"], dic["time"])
        dic["score"] = float(dic["score"])
    return results

def main():
    """Check the resulting csv file of packet classifications to see if actual attacks were detected"""
    parser = argparse.ArgumentParser()
    parser.add_argument("results_file", type=str, help="the results.csv file")
    parser.add_argument("attacks_file", type=str, help="the actual attacks file")
    parser.add_argument(
        "--thresh",
        type=str,
        default="0.5:0.5:1",
        help="range of thresholds to try. Format: start:stop:num_points, default: 0.5:0.5:1",
    )
    parser.add_argument("--plot", action="store_true", help="make plots")
    parser.add_argument(
        "--table",
        type=float,
        default=0.725,
        help="make table using the specified threshold",
    )
    args = parser.parse_args()

    check_results(
        args.results_file, args.attacks_file, args.thresh, args.plot, args.table
    )

if __name__ == "__main__":
    main()


import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import numpy as np

from utils.utils import ReplayBuffer, make_one_mini_batch, convert_to_tensor

class Agent(nn.Module):
    def __init__(self,algorithm, writer, device, state_dim, action_dim, args, demonstrations_location_args):
        super(Agent, self).__init__()
        self.writer = writer
        self.device = device
        self.args = args
        if self.args.on_policy == True :
            self.data = ReplayBuffer(action_prob_exist = True, max_size = self.args.traj_length, state_dim = state_dim, num_action = action_dim)
        else :
            self.data = ReplayBuffer(action_prob_exist = False, max_size = int(self.args.memory_size), state_dim = state_dim, num_action = action_dim)
        file_size = 120

        f = open(demonstrations_location_args.expert_state_location,'rb')
        self.expert_states = torch.tensor(np.concatenate([np.load(f) for _ in range(file_size)])).float()
        f = open(demonstrations_location_args.expert_action_location,'rb')
        self.expert_actions = torch.tensor(np.concatenate([np.load(f) for _ in range(file_size)]))
        f = open(demonstrations_location_args.expert_next_state_location,'rb')
        self.expert_next_states = torch.tensor(np.concatenate([np.load(f) for _ in range(file_size)])).float()
        f = open(demonstrations_location_args.expert_done_location,'rb')
        self.expert_dones = torch.tensor(np.concatenate([np.load(f) for _ in range(file_size)])).float().unsqueeze(-1)
        f.close()

        self.brain = algorithm

    def get_action(self,x):
        action, log_prob = self.brain.get_action(x)
        return action, log_prob

    def put_data(self,transition):
        self.data.put_data(transition)

    def train(self, discriminator, discriminator_batch_size, state_rms, n_epi, batch_size = 64):
        if self.args.on_policy :
            data = self.data.sample(shuffle = False)
            states, actions, rewards, next_states, done_masks, old_log_probs = convert_to_tensor(self.device, data['state'], data['action'], data['reward'], data['next_state'], data['done'], data['log_prob'])
        else :
            data = self.data.sample(shuffle = True, batch_size = discriminator_batch_size)
            states, actions, rewards, next_states, done_masks = convert_to_tensor(self.device, data['state'], data['action'], data['reward'], data['next_state'], data['done'])
        if discriminator.name() == 'sqil':
            agent_s,agent_a,agent_next_s,agent_done_mask = make_one_mini_batch(batch_size, states, actions, next_states, done_masks)
            expert_s,expert_a,expert_next_s,expert_done = make_one_mini_batch(batch_size, self.expert_states, self.expert_actions, self.expert_next_states, self.expert_dones)
            expert_done_mask = (1 - expert_done.float())

            discriminator.train_network(self.brain, n_epi, agent_s, agent_a, agent_next_s, agent_done_mask, expert_s, expert_a, expert_next_s, expert_done_mask)
            return
        if discriminator.args.is_airl == False:
            agent_s,agent_a = make_one_mini_batch(discriminator_batch_size, states, actions)
            expert_s,expert_a = make_one_mini_batch(discriminator_batch_size, self.expert_states, self.expert_actions)
            if self.args.on_policy :
                expert_s = np.clip((expert_s - state_rms.mean) / (state_rms.var ** 0.5 + 1e-8), -5, 5)
            discriminator.train_network(self.writer, n_epi, agent_s, agent_a, expert_s, expert_a)
        else:
            agent_s,agent_a,agent_next_s,agent_done_mask = make_one_mini_batch(discriminator_batch_size, states, actions, next_states, done_masks)
            expert_s,expert_a,expert_next_s,expert_done = make_one_mini_batch(discriminator_batch_size, self.expert_states, self.expert_actions, self.expert_next_states, self.expert_dones)

            expert_done_mask = (1 - expert_done.float())
            if self.args.on_policy :
                expert_s = np.clip((expert_s - state_rms.mean) / (state_rms.var ** 0.5 + 1e-8), -5, 5).float()
                expert_next_s = np.clip((expert_next_s - state_rms.mean) / (state_rms.var ** 0.5 + 1e-8), -5, 5).float()

            mu,sigma = self.brain.get_dist(agent_s.float().to(self.device))
            dist = torch.distributions.Normal(mu,sigma)
            agent_log_prob = dist.log_prob(agent_a).sum(-1,keepdim=True).detach()
            mu,sigma = self.brain.get_dist(expert_s.float().to(self.device))
            dist = torch.distributions.Normal(mu,sigma)
            expert_log_prob = dist.log_prob(expert_a).sum(-1,keepdim=True).detach()
            discriminator.train_network(self.writer, n_epi, agent_s, agent_a, agent_next_s,\
                                          agent_log_prob, agent_done_mask, expert_s, expert_a, expert_next_s, expert_log_prob, expert_done_mask)
        if self.args.on_policy :
            self.brain.train_network(self.writer, n_epi, states, actions, rewards, next_states, done_masks, old_log_probs)
        else :
            data = self.data.sample(shuffle = True, batch_size = batch_size)
            states, actions, rewards, next_states, done_masks = convert_to_tensor(self.device, data['state'], data['action'], data['reward'], data['next_state'], data['done'])
            self.brain.train_network(self.writer, n_epi, states, actions, rewards, next_states, done_masks)