# Copyright 2019 Nathan Jay and Noga Rotman
#
# 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.
import heapq
import os
import random
import sys
import time
import warnings
warnings.simplefilter(action='ignore', category=UserWarning)

import gym
import numpy as np
import pandas as pd
from gym import spaces
from gym.envs.registration import register
from gym.utils import seeding

from common import sender_obs
from common.utils import pcc_aurora_reward
from simulator.constants import (BYTES_PER_PACKET, EVENT_TYPE_ACK,
                                 EVENT_TYPE_SEND, MAX_RATE, MI_RTT_PROPORTION,
                                 MIN_RATE, REWARD_SCALE)
from simulator.link import Link
from simulator.trace import generate_traces

LATENCY_PENALTY = 1.0
LOSS_PENALTY = 1.0

USE_LATENCY_NOISE = True
MAX_LATENCY_NOISE = 1.1

# DEBUG = True
DEBUG = False

def debug_print(msg):
    if DEBUG:
        print(msg, file=sys.stderr, flush=True)

class EmuReplay:
    def __init__(self, ):
        df = pd.read_csv('aurora_emulation_log.csv')
        self.ts = df['timestamp'].tolist()
        self.send_rate = df['send_rate'].tolist()
        self.idx = 0

    def get_ts(self):
        if self.idx > len(self.ts):
            self.idx = len(self.ts) - 1
        ts = self.ts[self.idx]
        self.idx += 1
        return ts

    def get_rate(self):
        return self.send_rate[self.idx] / 8 / BYTES_PER_PACKET

    def reset(self):
        self.idx = 0

class Network():

    def __init__(self, senders, links, env):
        self.event_count = 0
        self.q = []
        self.cur_time = 0.0
        self.senders = senders
        self.links = links
        self.queue_initial_packets()
        self.env = env

        self.pkt_log = []

        self.recv_rate_cache = []

    def queue_initial_packets(self):
        for sender in self.senders:
            sender.register_network(self)
            sender.reset_obs()
            heapq.heappush(self.q, (0, sender, EVENT_TYPE_SEND,
                                    0, 0.0, False, self.event_count, sender.rto, 0))
            self.event_count += 1

    def reset(self):
        self.pkt_log = []
        self.cur_time = 0.0
        self.q = []
        [link.reset() for link in self.links]
        [sender.reset() for sender in self.senders]
        self.queue_initial_packets()
        self.recv_rate_cache = []

    def get_cur_time(self):
        return self.cur_time

    def run_for_dur(self, dur, action=None):
        if self.senders[0].lat_diff != 0:
            self.senders[0].start_stage = False
        start_time = self.cur_time
        end_time = min(self.cur_time + dur,
                       self.env.current_trace.timestamps[-1])
        debug_print('MI from {} to {}, dur {}'.format(
            self.cur_time, end_time, dur))
        for sender in self.senders:
            sender.reset_obs()
        # set_obs_start = False
        extra_delays = []  # time used to put packet onto the network
        while True:
            event_time, sender, event_type, next_hop, cur_latency, dropped, \
                event_id, rto, event_queue_delay = self.q[0]
            # if not sender.got_data and event_time >= end_time and event_type == EVENT_TYPE_ACK and next_hop == len(sender.path):
            #     end_time = event_time
            #     self.cur_time = end_time
            #     self.env.run_dur = end_time - start_time
            #     break
            if sender.got_data and event_time >= end_time and event_type == EVENT_TYPE_SEND:
                end_time = event_time
                self.cur_time = end_time
                break
            event_time, sender, event_type, next_hop, cur_latency, dropped, \
                event_id, rto, event_queue_delay = heapq.heappop(self.q)
            self.cur_time = event_time
            new_event_time = event_time
            new_event_type = event_type
            new_next_hop = next_hop
            new_latency = cur_latency
            new_dropped = dropped
            new_event_queue_delay = event_queue_delay
            push_new_event = False
            debug_print("Got %d event %s, to link %d, latency %f at time %f, "
                        "next_hop %d, dropped %s, event_q length %f, "
                        "sender rate %f, duration: %f, queue_size: %f, "
                        "rto: %f, cwnd: %f, ssthresh: %f, sender rto %f, "
                        "pkt in flight %d, wait time %d" % (
                            event_id, event_type, next_hop, cur_latency,
                            event_time, next_hop, dropped, len(self.q),
                            sender.rate, dur, self.links[0].queue_size,
                            rto, sender.cwnd, sender.ssthresh, sender.rto,
                            int(sender.bytes_in_flight/BYTES_PER_PACKET),
                            sender.pkt_loss_wait_time))
            if event_type == EVENT_TYPE_ACK:
                if next_hop == len(sender.path):
                    # if cur_latency > 1.0:
                    #     sender.timeout(cur_latency)
                    # sender.on_packet_lost(cur_latency)
                    if rto >= 0 and cur_latency > rto and sender.pkt_loss_wait_time <= 0:
                        sender.timeout()
                        dropped = True
                        new_dropped = True
                    elif dropped:
                        sender.on_packet_lost(cur_latency)
                        if not self.env.train_flag:
                            self.pkt_log.append(
                                [self.cur_time, event_id, 'lost',
                                 BYTES_PER_PACKET, cur_latency, event_queue_delay,
                                 self.links[0].pkt_in_queue,
                                 sender.rate * BYTES_PER_PACKET * 8,
                                 self.links[0].get_bandwidth(self.cur_time) * BYTES_PER_PACKET * 8])
                    else:
                        sender.on_packet_acked(cur_latency)
                        debug_print('Ack packet at {}'.format(self.cur_time))
                        # log packet acked
                        if not self.env.train_flag:
                            self.pkt_log.append(
                                [self.cur_time, event_id, 'acked',
                                 BYTES_PER_PACKET, cur_latency,
                                 event_queue_delay, self.links[0].pkt_in_queue,
                                 sender.rate * BYTES_PER_PACKET * 8,
                                 self.links[0].get_bandwidth(self.cur_time) * BYTES_PER_PACKET * 8])
                else:
                    if not self.env.train_flag:
                        self.pkt_log.append(
                            [self.cur_time, event_id, 'arrived',
                             BYTES_PER_PACKET, cur_latency, event_queue_delay,
                             self.links[0].pkt_in_queue,
                             sender.rate * BYTES_PER_PACKET * 8,
                             self.links[0].get_bandwidth(self.cur_time) * BYTES_PER_PACKET * 8])
                    new_next_hop = next_hop + 1
                    new_event_queue_delay += sender.path[next_hop].get_cur_queue_delay(
                        self.cur_time)
                    link_latency = sender.path[next_hop].get_cur_latency(
                        self.cur_time)
                    # link_latency *= self.env.current_trace.get_delay_noise_replay(self.cur_time)
                    # if USE_LATENCY_NOISE:
                    # link_latency *= random.uniform(1.0, MAX_LATENCY_NOISE)
                    new_latency += link_latency
                    new_event_time += link_latency
                    push_new_event = True
            elif event_type == EVENT_TYPE_SEND:
                if next_hop == 0:
                    if sender.can_send_packet():
                        sender.on_packet_sent()
                        # print('Send packet at {}'.format(self.cur_time))
                        if not self.env.train_flag:
                            self.pkt_log.append(
                                [self.cur_time, event_id, 'sent',
                                 BYTES_PER_PACKET, cur_latency,
                                 event_queue_delay, self.links[0].pkt_in_queue,
                                 sender.rate * BYTES_PER_PACKET * 8,
                                 self.links[0].get_bandwidth(self.cur_time) * BYTES_PER_PACKET * 8])
                        push_new_event = True
                    heapq.heappush(self.q, (self.cur_time + (1.0 / sender.rate),
                                            sender, EVENT_TYPE_SEND, 0, 0.0,
                                            False, self.event_count, sender.rto,
                                            0))
                    self.event_count += 1

                else:
                    push_new_event = True

                if next_hop == sender.dest:
                    new_event_type = EVENT_TYPE_ACK
                new_next_hop = next_hop + 1

                new_event_queue_delay += sender.path[next_hop].get_cur_queue_delay(
                    self.cur_time)
                link_latency = sender.path[next_hop].get_cur_latency(
                    self.cur_time)
                # if USE_LATENCY_NOISE:
                # link_latency *= random.uniform(1.0, MAX_LATENCY_NOISE)
                # link_latency += self.env.current_trace.get_delay_noise(
                #     self.cur_time, self.links[0].get_bandwidth(self.cur_time)) / 1000
                # link_latency += max(0, np.random.normal(0, 1) / 1000)
                # link_latency += max(0, np.random.uniform(0, 5) / 1000)
                # link_latency *= self.env.current_trace.get_delay_noise_replay(self.cur_time)
                new_latency += link_latency
                new_event_time += link_latency
                new_dropped = not sender.path[next_hop].packet_enters_link(
                    self.cur_time)
                extra_delays.append(
                    1 / self.links[0].get_bandwidth(self.cur_time))
                if not new_dropped:
                    sender.queue_delay_samples.append(new_event_queue_delay)

            if push_new_event:
                heapq.heappush(self.q, (new_event_time, sender, new_event_type,
                                        new_next_hop, new_latency, new_dropped,
                                        event_id, rto, new_event_queue_delay))
        for sender in self.senders:
            sender.record_run()

        sender_mi = self.senders[0].history.back() #get_run_data()
        throughput = sender_mi.get("recv rate")  # bits/sec
        latency = sender_mi.get("avg latency")  # second
        loss = sender_mi.get("loss ratio")
        debug_print("thpt %f, delay %f, loss %f, bytes sent %f, bytes acked %f" % (
            throughput/1e6, latency, loss, sender_mi.bytes_sent, sender_mi.bytes_acked))
        reward = pcc_aurora_reward(
            throughput / 8 / BYTES_PER_PACKET, latency, loss,
            np.mean(self.env.current_trace.bandwidths) * 1e6 / 8 / BYTES_PER_PACKET,
            np.mean(self.env.current_trace.delays) * 2 / 1e3)

        if latency > 0.0:
            self.env.run_dur = MI_RTT_PROPORTION * \
                sender_mi.get("avg latency") + np.mean(extra_delays)
        # elif self.env.run_dur != 0.01:
            # assert self.env.run_dur >= 0.03
            # self.env.run_dur = max(MI_RTT_PROPORTION * sender_mi.get("avg latency"), 5 * (1 / self.senders[0].rate))

        self.senders[0].avg_latency = sender_mi.get("avg latency")  # second
        self.senders[0].recv_rate = round(sender_mi.get("recv rate"), 3)  # bits/sec
        self.senders[0].send_rate = round(sender_mi.get("send rate"), 3)  # bits/sec
        self.senders[0].lat_diff = sender_mi.rtt_samples[-1] - sender_mi.rtt_samples[0]
        self.senders[0].latest_rtt = sender_mi.rtt_samples[-1]
        self.recv_rate_cache.append(self.senders[0].recv_rate)
        if len(self.recv_rate_cache) > 6:
            self.recv_rate_cache = self.recv_rate_cache[1:]
        self.senders[0].max_tput = max(self.recv_rate_cache)

        if self.senders[0].lat_diff == 0 and self.senders[0].start_stage:  # no latency change
            pass
            # self.senders[0].max_tput = max(self.senders[0].recv_rate, self.senders[0].max_tput)
        elif self.senders[0].lat_diff == 0 and not self.senders[0].start_stage:  # no latency change
            pass
            # self.senders[0].max_tput = max(self.senders[0].recv_rate, self.senders[0].max_tput)
        elif self.senders[0].lat_diff > 0:  # latency increase
            self.senders[0].start_stage = False
            # self.senders[0].max_tput = self.senders[0].recv_rate # , self.max_tput)
        else:  # latency decrease
            self.senders[0].start_stage = False
            # self.senders[0].max_tput = max(self.senders[0].recv_rate, self.senders[0].max_tput)
        return reward * REWARD_SCALE

class Sender():

    def __init__(self, rate, path, dest, features, cwnd=25, history_len=10,
                 delta_scale=1):
        self.id = Sender._get_next_id()
        self.delta_scale = delta_scale
        self.starting_rate = rate
        self.rate = rate
        self.sent = 0
        self.acked = 0
        self.lost = 0
        self.bytes_in_flight = 0
        self.min_latency = None
        self.rtt_samples = []
        self.rtt_samples_ts = []
        self.queue_delay_samples = []
        self.prev_rtt_samples = self.rtt_samples
        self.sample_time = []
        self.net = None
        self.path = path
        self.dest = dest
        self.history_len = history_len
        self.features = features
        self.history = sender_obs.SenderHistory(self.history_len,
                                                self.features, self.id)
        self.cwnd = cwnd
        self.use_cwnd = False
        self.rto = -1
        self.ssthresh = 0
        self.pkt_loss_wait_time = -1
        self.estRTT = 1000000 / 1e6  # SynInterval in emulation
        self.RTTVar = self.estRTT / 2  # RTT variance
        self.got_data = False

        self.min_rtt = 10
        self.max_tput = 0
        self.start_stage = True
        self.lat_diff = 0
        self.recv_rate = 0
        self.send_rate = 0
        self.avg_latency = 0
        self.latest_rtt = 0

    _next_id = 1

    def _get_next_id():
        result = Sender._next_id
        Sender._next_id += 1
        return result

    def apply_rate_delta(self, delta):
        # if self.got_data:
        delta *= self.delta_scale
        #print("Applying delta %f" % delta)
        if delta >= 0.0:
            self.set_rate(self.rate * (1.0 + delta))
        else:
            self.set_rate(self.rate / (1.0 - delta))

    def apply_cwnd_delta(self, delta):
        delta *= self.delta_scale
        #print("Applying delta %f" % delta)
        if delta >= 0.0:
            self.set_cwnd(self.cwnd * (1.0 + delta))
        else:
            self.set_cwnd(self.cwnd / (1.0 - delta))

    def can_send_packet(self):
        if self.use_cwnd:
            return int(self.bytes_in_flight) / BYTES_PER_PACKET < self.cwnd
        else:
            return True

    def register_network(self, net):
        self.net = net

    def on_packet_sent(self):
        self.sent += 1
        self.bytes_in_flight += BYTES_PER_PACKET

    def on_packet_acked(self, rtt):
        self.min_rtt = min(self.min_rtt, rtt)
        self.estRTT = (7.0 * self.estRTT + rtt) / 8.0  # RTT of emulation way
        self.RTTVar = (self.RTTVar * 7.0 + abs(rtt - self.estRTT) * 1.0) / 8.0

        self.acked += 1
        self.rtt_samples.append(rtt)
        self.rtt_samples_ts.append(self.net.get_cur_time())
        # self.rtt_samples.append(self.estRTT)
        if (self.min_latency is None) or (rtt < self.min_latency):
            self.min_latency = rtt
        self.bytes_in_flight -= BYTES_PER_PACKET
        if not self.got_data:
            self.got_data = len(self.rtt_samples) >= 1
        # self.got_data = True

    def on_packet_lost(self, rtt):
        self.lost += 1
        self.bytes_in_flight -= BYTES_PER_PACKET

    def set_rate(self, new_rate):
        self.rate = new_rate
        # print("Attempt to set new rate to %f (min %f, max %f)" % (new_rate, MIN_RATE, MAX_RATE))
        if self.rate > MAX_RATE:
            self.rate = MAX_RATE
        if self.rate < MIN_RATE:
            self.rate = MIN_RATE

    def set_cwnd(self, new_cwnd):
        self.cwnd = int(new_cwnd)
        #print("Attempt to set new rate to %f (min %f, max %f)" % (new_rate, MIN_RATE, MAX_RATE))
        # if self.cwnd > MAX_CWND:
        #     self.cwnd = MAX_CWND
        # if self.cwnd < MIN_CWND:
        #     self.cwnd = MIN_CWND

    def record_run(self):
        smi = self.get_run_data()
        # if not self.got_data and smi.rtt_samples:
        #     self.got_data = True
        #     self.history.step(smi)
        # else:
        self.history.step(smi)

    def get_obs(self):
        return self.history.as_array()

    def get_run_data(self):
        obs_end_time = self.net.get_cur_time()

        #obs_dur = obs_end_time - self.obs_start_time
        #print("Got %d acks in %f seconds" % (self.acked, obs_dur))
        #print("Sent %d packets in %f seconds" % (self.sent, obs_dur))
        #print("self.rate = %f" % self.rate)
        # print(self.acked, self.sent)
        if not self.rtt_samples and self.prev_rtt_samples:
            rtt_samples = [np.mean(self.prev_rtt_samples)]
        else:
            rtt_samples = self.rtt_samples
        # if not self.rtt_samples:
        #     print(self.obs_start_time, obs_end_time, self.rate)
        # rtt_samples is empty when there is no packet acked in MI
        # Solution: inherit from previous rtt_samples.

        # recv_start = self.rtt_samples_ts[0] if len(
        #     self.rtt_samples) >= 2 else self.obs_start_time
        recv_start = self.history.back().recv_end if len(
            self.rtt_samples) >= 1 else self.obs_start_time
        recv_end = self.rtt_samples_ts[-1] if len(
            self.rtt_samples) >= 1 else obs_end_time
        bytes_acked = self.acked * BYTES_PER_PACKET
        if recv_start == 0:
            recv_start = self.rtt_samples_ts[0]
            bytes_acked = (self.acked - 1) * BYTES_PER_PACKET

        # bytes_acked = max(0, (self.acked-1)) * BYTES_PER_PACKET if len(
        #     self.rtt_samples) >= 2 else self.acked * BYTES_PER_PACKET
        return sender_obs.SenderMonitorInterval(
            self.id,
            bytes_sent=self.sent * BYTES_PER_PACKET,
            # max(0, (self.acked-1)) * BYTES_PER_PACKET,
            # bytes_acked=self.acked * BYTES_PER_PACKET,
            bytes_acked=bytes_acked,
            bytes_lost=self.lost * BYTES_PER_PACKET,
            send_start=self.obs_start_time,
            send_end=obs_end_time,
            # recv_start=self.obs_start_time,
            # recv_end=obs_end_time,
            recv_start=recv_start,
            recv_end=recv_end,
            rtt_samples=rtt_samples,
            queue_delay_samples=self.queue_delay_samples,
            packet_size=BYTES_PER_PACKET
        )

    def reset_obs(self):
        self.sent = 0
        self.acked = 0
        self.lost = 0
        if self.rtt_samples:
            self.prev_rtt_samples = self.rtt_samples
        self.rtt_samples = []
        self.rtt_samples_ts = []
        self.queue_delay_samples = []
        self.obs_start_time = self.net.get_cur_time()

    def print_debug(self):
        print("Sender:")
        print("Obs: %s" % str(self.get_obs()))
        print("Rate: %f" % self.rate)
        print("Sent: %d" % self.sent)
        print("Acked: %d" % self.acked)
        print("Lost: %d" % self.lost)
        print("Min Latency: %s" % str(self.min_latency))

    def reset(self):
        #print("Resetting sender!")
        self.rate = self.starting_rate
        self.bytes_in_flight = 0
        self.min_latency = None
        self.reset_obs()
        self.history = sender_obs.SenderHistory(self.history_len,
                                                self.features, self.id)
        self.estRTT = 1000000 / 1e6  # SynInterval in emulation
        self.RTTVar = self.estRTT / 2  # RTT variance

        self.got_data = False
        self.min_rtt = 10
        self.max_tput = 0
        self.start_stage = True
        self.lat_diff = 0
        self.recv_rate = 0
        self.send_rate = 0
        self.avg_latency = 0
        self.latest_rtt = 0

    def timeout(self):
        # placeholder
        pass

class SimulatedNetworkEnv(gym.Env):

    def __init__(self, traces, history_len=10,
                 # features="sent latency inflation,latency ratio,send ratio",
                 features="sent latency inflation,latency ratio,recv ratio",
                 congestion_control_type="aurora", train_flag=False,
                 delta_scale=1.0, config_file=None):
        """Network environment used in simulation.
        congestion_control_type: aurora is pcc-rl. cubic is TCPCubic.
        """
        assert congestion_control_type in {"aurora", "cubic"}, \
            "Unrecognized congestion_control_type {}.".format(
                congestion_control_type)
        # self.replay = EmuReplay()
        self.config_file = config_file
        self.delta_scale = delta_scale
        self.traces = traces
        self.current_trace = np.random.choice(self.traces)
        self.train_flag = train_flag
        self.congestion_control_type = congestion_control_type
        if self.congestion_control_type == 'aurora':
            self.use_cwnd = False
        elif self.congestion_control_type == 'cubic':
            self.use_cwnd = True

        self.history_len = history_len
        # print("History length: %d" % history_len)
        self.features = features.split(",")
        # print("Features: %s" % str(self.features))

        self.links = None
        self.senders = None
        self.create_new_links_and_senders()
        self.net = Network(self.senders, self.links, self)
        self.run_dur = None
        self.run_period = 0.1
        self.steps_taken = 0
        self.debug_thpt_changes = False
        self.last_thpt = None
        self.last_rate = None

        if self.use_cwnd:
            self.action_space = spaces.Box(
                np.array([-1e12, -1e12]), np.array([1e12, 1e12]), dtype=np.float32)
        else:
            self.action_space = spaces.Box(
                np.array([-1e12]), np.array([1e12]), dtype=np.float32)

        self.observation_space = None
        # use_only_scale_free = True
        single_obs_min_vec = sender_obs.get_min_obs_vector(self.features)
        single_obs_max_vec = sender_obs.get_max_obs_vector(self.features)
        self.observation_space = spaces.Box(np.tile(single_obs_min_vec, self.history_len),
                                            np.tile(single_obs_max_vec,
                                                    self.history_len),
                                            dtype=np.float32)
        # single_obs_min_vec = np.array([0, 0, -1e12, 0, 0, 0, 0])
        # single_obs_max_vec =  np.array([1e12, 1e12, 1e12, 1, 1e12, 1e12, 1e12])
        # self.observation_space = spaces.Box(single_obs_min_vec,
        #                                     single_obs_max_vec,
        #                                     dtype=np.float32)

        self.reward_sum = 0.0
        self.reward_ewma = 0.0

        self.episodes_run = -1

    def seed(self, seed=None):
        self.rand, seed = seeding.np_random(seed)
        return [seed]

    def _get_all_sender_obs(self):
        sender_obs = self.senders[0].get_obs()
        sender_obs = np.array(sender_obs).reshape(-1,)
        return sender_obs

    def step(self, actions):
        #print("Actions: %s" % str(actions))
        # print(actions)
        for i in range(0, 1):  # len(actions)):
            #print("Updating rate for sender %d" % i)
            action = actions
            self.senders[i].apply_rate_delta(action[0])
            if self.use_cwnd:
                self.senders[i].apply_cwnd_delta(action[1])
        # print("Running for %fs" % self.run_dur)
        reward = self.net.run_for_dur(self.run_dur, action=actions[0])
        self.steps_taken += 1
        sender_obs = self._get_all_sender_obs()

        should_stop = self.current_trace.is_finished(self.net.get_cur_time())

        self.reward_sum += reward
        # print('env step: {}s'.format(time.time() - t_start))

        # sender_obs = np.array([self.senders[0].send_rate,
        #         self.senders[0].avg_latency,
        #         self.senders[0].lat_diff, int(self.senders[0].start_stage),
        #         self.senders[0].max_tput, self.senders[0].min_rtt,
        #         self.senders[0].latest_rtt])
        return sender_obs, reward, should_stop, {}

    def print_debug(self):
        print("---Link Debug---")
        for link in self.links:
            link.print_debug()
        print("---Sender Debug---")
        for sender in self.senders:
            sender.print_debug()

    def create_new_links_and_senders(self):
        # self.replay.reset()
        self.links = [Link(self.current_trace), Link(self.current_trace)]
        if self.congestion_control_type == "aurora":
            if not self.train_flag:

                self.senders = [Sender(  # self.replay.get_rate(),
                    # 2500000 / 8 /BYTES_PER_PACKET / 0.048,
                    # 12000000 / 8 /BYTES_PER_PACKET / 0.048,
                    10 / (self.current_trace.get_delay(0) *2/1000),
                    # 100,
                    [self.links[0], self.links[1]], 0,
                    self.features,
                    history_len=self.history_len,
                    delta_scale=self.delta_scale)]
            else:
                # self.senders = [Sender(random.uniform(0.3, 1.5) * bw,
                #                        [self.links[0], self.links[1]], 0,
                #                        self.features,
                #                        history_len=self.history_len)]
                # self.senders = [Sender(random.uniform(10/bw, 1.5) * bw,
                #                        [self.links[0], self.links[1]], 0,
                #                        self.features,
                #                        history_len=self.history_len,
                #                        delta_scale=self.delta_scale)]
                self.senders = [Sender(
                    # 100,
                    10 / (self.current_trace.get_delay(0) *2/1000),
                                       [self.links[0], self.links[1]], 0,
                                       self.features,
                                       history_len=self.history_len,
                                       delta_scale=self.delta_scale)]
        elif self.congestion_control_type == "cubic":
            raise NotImplementedError
        else:
            raise RuntimeError("Unrecognized congestion_control_type {}".format(
                self.congestion_control_type))
        # self.run_dur = 3 * lat
        # self.run_dur = 1 * lat
        if not self.senders[0].rtt_samples:
            # self.run_dur = 0.473
            # self.run_dur = 5 / self.senders[0].rate
            self.run_dur = 0.01
            # self.run_dur = self.current_trace.get_delay(0) * 2 / 1000
            # self.run_dur = self.replay.get_ts() -  0

    def reset(self):
        self.steps_taken = 0
        self.net.reset()
        self.current_trace = np.random.choice(self.traces)
        self.current_trace.reset()
        self.create_new_links_and_senders()
        self.net = Network(self.senders, self.links, self)
        self.episodes_run += 1
        if self.train_flag and self.config_file is not None and self.episodes_run % 100 == 0:
            print('change traces', self.episodes_run)
            self.traces = generate_traces(self.config_file, 10,
                                          duration=10, constant_bw=False)
        # self.replay.reset()
        self.net.run_for_dur(self.run_dur)
        self.reward_ewma *= 0.99
        self.reward_ewma += 0.01 * self.reward_sum
        # print("Reward: %0.2f, Ewma Reward: %0.2f" % (self.reward_sum, self.reward_ewma))
        self.reward_sum = 0.0
        return self._get_all_sender_obs()
        # return np.array([self.senders[0].send_rate, self.senders[0].avg_latency,
        #         self.senders[0].lat_diff, int(self.senders[0].start_stage),
        #         self.senders[0].max_tput, self.senders[0].min_rtt,
        #         self.senders[0].latest_rtt])

register(id='PccNs-v0', entry_point='simulator.network:SimulatedNetworkEnv')


# Copyright (c) 2020 NVIDIA Corporation. All rights reserved.
# This work is licensed under the NVIDIA Source Code License - Non-commercial. Full
# text can be found in LICENSE.md

from __future__ import division
import numpy as np
import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data
from torch.autograd import Variable
import os
import datetime
import re
import matplotlib.pyplot as plt
import time
from transforms3d.quaternions import *
from transforms3d.euler import *
from transforms3d.axangles import *
import scipy
from ycb_render.ycb_renderer import *
from decimal import *
import cv2
from shutil import copyfile
from networks.aae_models import *
from config.config import cfg
import matplotlib.patches as patches
from mpl_toolkits.mplot3d import axes3d, Axes3D
import gc

class aae_trainer(nn.Module):
    def __init__(self, cfg_path, object_names, modality, config_new=None,
                 aae_capacity=1, aae_code_dim=128, ckpt_path=None, obj_ctg='ycb', lr=0.0002):
        super(aae_trainer, self).__init__()

        self.cfg_path = cfg_path

        if config_new != None:
            self.cfg_all = config_new
        else:
            self.cfg_all = cfg

        self.obj_ctg = obj_ctg
        self.modality = modality

        if not os.path.exists('./checkpoints'):
            os.mkdir('./checkpoints')
        self.ckpt_dir = './checkpoints'

        self.AAE = AAE(object_names=object_names,
                       modality=modality,
                       capacity=aae_capacity,
                       code_dim=aae_code_dim,
                       model_path=ckpt_path)

        self.object_names = object_names

        self.use_GPU = (torch.cuda.device_count() > 0)

        self.code_dim = aae_code_dim

        if self.modality == 'rgbd':
            self.optimizer = optim.Adam(list(self.AAE.encoder.parameters()) + \
                                        list(self.AAE.decoder.parameters()) + \
                                        list(self.AAE.depth_decoder.parameters()),
                                        lr=lr)
        else:
            self.optimizer = optim.Adam(list(self.AAE.encoder.parameters()) + \
                                        list(self.AAE.decoder.parameters()),
                                        lr=lr)

        self.mseloss = nn.MSELoss()
        self.l1_loss = nn.L1Loss()
        self.l1_recon_loss = nn.L1Loss(reduction='mean')

        if self.use_GPU:
            self.mseloss = self.mseloss.cuda()
            self.l1_loss = self.l1_loss.cuda()
            self.l1_recon_loss = self.l1_recon_loss.cuda()

        self.loss_history_recon = []
        self.val_loss_history_recon = []

        self.batch_size_train = self.cfg_all.TRAIN.BATCH_SIZE
        self.batch_size_val = self.cfg_all.TRAIN.VAL_BATCH_SIZE
        self.start_epoch = 1

        self.codebook_dir = None
        self.log_dir = None
        self.checkpoint_path = None

        self.lb_shift = self.cfg_all.TRAIN.SHIFT_MIN
        self.ub_shift = self.cfg_all.TRAIN.SHIFT_MAX
        self.lb_scale = self.cfg_all.TRAIN.SCALE_MIN
        self.ub_scale = self.cfg_all.TRAIN.SCALE_MAX

        if ckpt_path is not None:
            self.load_ckpt(ckpt_path=ckpt_path)

    def set_log_dir(self, dataset_name='', model_path=None, now=None, ):
        # Set date and epoch counter as if starting a new model
        self.epoch = 0
        if now == None:
            now = datetime.datetime.now()

        # If we have a model path with date and epochs use them
        if model_path:
            regex = r".*/\w+(\d{4})(\d{2})(\d{2})T(\d{2})(\d{2})(\d{2})/trans\_\w+(\d{4})\.pth"
            m = re.match(regex, model_path)
            if m:
                now = datetime.datetime(int(m.group(1)), int(m.group(2)), int(m.group(3)),
                                        int(m.group(4)), int(m.group(5)), int(m.group(6)))

        # Directory for training logs
        self.log_dir = os.path.join(self.ckpt_dir, "{}_{}_{}_{:%Y%m%dT%H%M%S}".format(
            dataset_name, self.object_names[0], self.cfg_all.EXP_NAME, now))
        # Path to save after each epoch. Include placeholders that get filled by Keras.
        self.checkpoint_path = os.path.join(self.log_dir, "ckpt_{}_*epoch*.pth".format(
            self.obj_ctg))
        self.checkpoint_path = self.checkpoint_path.replace(
            "*epoch*", "{:04d}")

    def save_ckpt(self, epoch):
        print('=> Saving checkpoint to {} ...'.format(self.checkpoint_path.format(epoch)))
        torch.save({
            'epoch': epoch,
            'log_dir': self.log_dir,
            'checkpoint_path': self.checkpoint_path,
            'aae_state_dict': self.AAE.state_dict(),
            'optimizer': self.optimizer.state_dict(),
            'loss_history_recon': self.loss_history_recon,
            'val_loss_history_recon': self.val_loss_history_recon,
        }, self.checkpoint_path.format(epoch))
        print('=> Finished saving checkpoint to {} ! '.format(self.checkpoint_path.format(epoch)))

    def load_ckpt(self, ckpt_path):
        if os.path.isfile(ckpt_path):
            print("=> Loading checkpoint from {} ...".format(ckpt_path))
            checkpoint = torch.load(ckpt_path)
            self.start_epoch = checkpoint['epoch'] + 1
            self.log_dir = checkpoint['log_dir']
            self.checkpoint_path = checkpoint['checkpoint_path']
            self.AAE.load_ckpt_weights(checkpoint['aae_state_dict'])
            self.optimizer.load_state_dict(checkpoint['optimizer'])
            self.loss_history_recon = checkpoint['loss_history_recon']
            self.val_loss_history_recon = checkpoint['val_loss_history_recon']
            print("=> Finished loading checkpoint from {} (epoch {})"
                  .format(ckpt_path, checkpoint['epoch']))
        else:
            print('=> Cannot find checkpoint file in {} !'.format(ckpt_path))

    def plot_loss(self, loss, title, save=True, log_dir=None):
        loss = np.array(loss)
        plt.figure(title)
        plt.gcf().clear()
        plt.plot(loss, label='train')
        plt.xlabel('epoch')
        plt.ylabel('loss')
        plt.legend()
        if save:
            save_path = os.path.join(log_dir, "{}.png".format(title))
            plt.savefig(save_path)
            plt.close()
        else:
            plt.show(block=False)
            plt.pause(0.1)

    def train_model(self, train_dataset, epochs, dstr_dataset=None, save_frequency=5):
        self.AAE.encoder.train()
        self.AAE.decoder.train()
        if self.modality == 'rgbd':
            self.AAE.depth_decoder.train()

        train_set = train_dataset
        print('train set size {}'.format(len(train_set)))

        if self.log_dir == None:
            self.set_log_dir(dataset_name=train_dataset._name)
            if not os.path.exists(self.log_dir) and save_frequency > 0:
                print('Create folder at {}'.format(self.log_dir))
                os.makedirs(self.log_dir)
                copyfile(self.cfg_path, self.log_dir + '/config.yml')

        print('dataset workers %d' % (self.cfg_all.TRAIN.WORKERS))
        train_generator = torch.utils.data.DataLoader(train_set, batch_size=self.batch_size_train,
                                                      shuffle=True, num_workers=self.cfg_all.TRAIN.WORKERS)

        if dstr_dataset != None:
            print('background workers %d' % (self.cfg_all.TRAIN.DISTRACTOR_WORKERS))
            train_dstr_generator = torch.utils.data.DataLoader(dstr_dataset, batch_size=self.batch_size_train,
                                                              shuffle=True, num_workers=self.cfg_all.TRAIN.DISTRACTOR_WORKERS)
        else:
            train_dstr_generator = None

        train_steps = np.floor(len(train_set)/self.batch_size_train)
        train_steps = np.floor(train_steps/4)

        for epoch in np.arange(start=self.start_epoch, stop=(self.start_epoch+epochs)):
            print("Epoch {}/{}.".format(epoch, (self.start_epoch+epochs)-1))

            if self.modality == 'rgbd':
                recon_loss_rgb, recon_loss_train = self.train_epoch_rgbd(train_generator, self.optimizer,
                                                                          train_steps, epoch, self.start_epoch + epochs - 1,
                                                                          dstrgenerator=train_dstr_generator)
                self.loss_history_recon.append(recon_loss_rgb + recon_loss_train)

            else:
                recon_loss_rgb, recon_loss_train = self.train_epoch_rgb(train_generator, self.optimizer,
                                        train_steps, epoch, self.start_epoch+epochs-1,
                                        dstrgenerator=train_dstr_generator)
                self.loss_history_recon.append(recon_loss_rgb + recon_loss_train)

            self.plot_loss(self.loss_history_recon, 'recon loss', save=True, log_dir=self.log_dir)
            if save_frequency > 0 and epoch % save_frequency == 0:
                self.save_ckpt(epoch)

    def train_epoch_rgbd(self, datagenerator, optimizer, steps, epoch, total_epoch, dstrgenerator=None):
        loss_sum_rgb = 0
        loss_sum_depth = 0
        step = 0
        optimizer.zero_grad()

        if dstrgenerator != None:
            enum_dstrgenerator = enumerate(dstrgenerator)

        for inputs in datagenerator:

            # receiving data from the renderer
            images, images_target, pose_cam, mask,\
            translation_target, scale_target, \
            affine_target, roi_center, roi_size, roi_affine, depths, depths_target = inputs

            if self.use_GPU:
                images = images.cuda()
                images_target = images_target.cuda()
                mask = mask.cuda()
                roi_affine = roi_affine.cuda()
                roi_center = roi_center.cuda().float()
                roi_size = roi_size.cuda().float()

            # warp the images according to the center and size of rois
            grids = F.affine_grid(roi_affine, images.size())
            images = F.grid_sample(images, grids)
            depths = F.grid_sample(depths, grids)
            mask = F.grid_sample(mask, grids)
            mask = 1 - mask

            # add random background and gaussian noise
            if dstrgenerator != None:
                _, images_dstr = next(enum_dstrgenerator)
                if images_dstr.size(0) != images.size(0):
                    enum_dstrgenerator = enumerate(dstrgenerator)
                    _, images_dstr = next(enum_dstrgenerator)
                if self.use_GPU:
                    images_dstr = images_dstr.cuda()
                images = images + mask * images_dstr
            noise_level = np.random.uniform(0, 0.05)
            images += torch.randn_like(images) * noise_level

            # add random background to depth image
            _, depth_dstr = next(enum_dstrgenerator)
            if depth_dstr.size(0) != depths.size(0):
                enum_dstrgenerator = enumerate(dstrgenerator)
                _, depth_dstr = next(enum_dstrgenerator)
            if self.use_GPU:
                depth_dstr = depth_dstr.cuda()
            depths_background = torch.sum(mask * depth_dstr, dim=1, keepdim=True) / np.random.uniform(0.5, 2.0)
            depths = depths + depths_background
            depths += torch.rand_like(depths) * np.random.uniform(0, 0.05)
            depths = torch.clamp(depths, 0, 1)

            # construct tensor for roi information
            roi_info = torch.zeros(images.size(0), 5).float().cuda()
            roi_center += torch.from_numpy(np.random.uniform(self.cfg_all.TRAIN.SHIFT_MIN,
                                                                   self.cfg_all.TRAIN.SHIFT_MAX,
                                                                   size=(roi_info.size(0), 2))).float().cuda()
            roi_size = roi_size * torch.from_numpy(np.random.uniform(self.cfg_all.TRAIN.SCALE_MIN,
                                                                     self.cfg_all.TRAIN.SCALE_MAX,
                                                                     size=(roi_info.size(0), ))).float().cuda()
            roi_info[:, 0] = torch.arange(images.size(0))
            roi_info[:, 1] = roi_center[:, 0] - roi_size / 2
            roi_info[:, 2] = roi_center[:, 1] - roi_size / 2
            roi_info[:, 3] = roi_center[:, 0] + roi_size / 2
            roi_info[:, 4] = roi_center[:, 1] + roi_size / 2

            roi_info_copy = roi_info.clone()
            roi_info_copy_depth = roi_info.clone()

            # # visualization for debugging
            # roi_info_copy2 = roi_info.clone()
            # images_roi = ROIAlign((128, 128), 1.0, 0)(images, roi_info_copy)
            # images_roi_disp = images_roi[0].permute(1, 2, 0).cpu().numpy()
            # depth_roi = ROIAlign((128, 128), 1.0, 0)(depths, roi_info_copy2)
            # depth_roi_disp = depth_roi[0, 0].cpu().numpy()
            # image_disp = images[0].permute(1, 2, 0).cpu().numpy()
            # depth_disp = depths[0, 0].cpu().numpy()
            # depth_target_disp = depths_target[0, 0].cpu().numpy()
            # mask_disp = mask[0].permute(1, 2, 0).repeat(1, 1, 3).cpu().numpy()
            # plt.figure()
            # plt.subplot(2, 3, 1)
            # plt.imshow(np.concatenate((image_disp, mask_disp), axis=1))
            # plt.subplot(2, 3, 2)
            # plt.imshow(images_roi_disp)
            # plt.subplot(2, 3, 3)
            # plt.imshow(images_target[0].permute(1, 2, 0).cpu().numpy())
            # plt.subplot(2, 3, 4)
            # plt.imshow(depth_disp)
            # plt.subplot(2, 3, 5)
            # plt.imshow(depth_roi_disp)
            # plt.subplot(2, 3, 6)
            # plt.imshow(depth_target_disp)
            # plt.show()

            # AAE forward pass
            images_input = torch.cat((images, depths), dim=1)
            outputs = self.AAE.forward_rgbd(images_input, roi_info)
            images_recnst = outputs[0]
            loss_reconstr = self.AAE.B_loss(images_recnst[:, :3, :, :], images_target.detach())
            loss_depth = self.AAE.B_loss(images_recnst[:, [3], :, :], depths_target)
            loss = loss_reconstr + loss_depth

            loss_aae_rgb_data = loss_reconstr.data.cpu().item()
            loss_aae_depth_data = loss_depth.data.cpu().item()

            # AAE backward pass
            optimizer.zero_grad()
            try:
                loss.backward()
            except:
                pass

            optimizer.step()

            print("{}/{}: {}/{}, loss_rgb: {:.4f}, loss depth: {:.4f}".format(epoch, total_epoch, step + 1, int(steps),
                  loss_aae_rgb_data, loss_aae_depth_data))

            loss_sum_rgb += loss_aae_rgb_data / steps
            loss_sum_depth += loss_aae_depth_data / steps

            # display
            plot_n_comparison = 20
            if step < plot_n_comparison:
                images_roi = ROIAlign((128, 128), 1.0, 0)(images, roi_info_copy)
                image = images_roi[0].detach().permute(1, 2, 0).cpu().numpy()
                image_target = images_target[0].permute(1, 2, 0).cpu().numpy()
                depths_roi = ROIAlign((128, 128), 1.0, 0)(depths, roi_info_copy_depth)
                depth = depths_roi[0, 0].detach().cpu().numpy()
                depth_target = depths_target[0, 0].cpu().numpy()
                depth_recon = images_recnst[0, 3].detach().cpu().numpy()
                image_recon = images_recnst[0, :3].detach().permute(1, 2, 0).cpu().numpy()
                disp = (image, image_target, image_recon, depth, depth_target, depth_recon)
                self.plot_comparison(disp, str(step))

            if step==steps-1:
                break
            step += 1

        return loss_sum_rgb, loss_sum_depth

    def train_epoch_rgb(self, datagenerator, optimizer, steps, epoch, total_epoch, dstrgenerator=None, visualize=False):
        loss_sum_rgb = 0
        loss_sum_depth = 0
        step = 0
        optimizer.zero_grad()

        if dstrgenerator != None:
            enum_dstrgenerator = enumerate(dstrgenerator)

        for inputs in datagenerator:

            # receiving data from the renderer
            images, images_target, pose_cam, mask,\
            translation_target, scale_target, \
            affine_target, roi_center, roi_size, roi_affine, depths, depths_target = inputs

            if self.use_GPU:
                images = images.cuda()
                images_target = images_target.cuda()
                mask = mask.cuda()
                roi_affine = roi_affine.cuda()
                roi_center = roi_center.cuda().float()
                roi_size = roi_size.cuda().float()

            # warp the images according to the center and size of rois
            grids = F.affine_grid(roi_affine, images.size())
            images = F.grid_sample(images, grids)
            depths = F.grid_sample(depths, grids)
            mask = F.grid_sample(mask, grids)
            mask = 1 - mask

            # add random background and gaussian noise
            if dstrgenerator != None:
                _, images_dstr = next(enum_dstrgenerator)
                if images_dstr.size(0) != images.size(0):
                    enum_dstrgenerator = enumerate(dstrgenerator)
                    _, images_dstr = next(enum_dstrgenerator)
                if self.use_GPU:
                    images_dstr = images_dstr.cuda()
                images = images + mask * images_dstr
            noise_level = np.random.uniform(0, 0.05)
            images += torch.randn_like(images) * noise_level

            class_info = torch.ones((images.size(0), 1, 128, 128), dtype=torch.float32).cuda()

            # visualization
            if visualize:
                image_disp = images[0].permute(1, 2, 0).cpu().numpy()
                image_target_disp = images_target[0].permute(1, 2, 0).cpu().numpy()
                depth_disp = depths[0, 0].cpu().numpy()
                depth_target_disp = depths_target[0, 0].cpu().numpy()
                mask_disp = mask[0].permute(1, 2, 0).repeat(1, 1, 3).cpu().numpy()
                plt.figure()
                plt.subplot(2, 2, 1)
                im = np.concatenate((image_disp, mask_disp), axis=1)
                im = np.clip(im * 255, 0, 255).astype(np.uint8)
                plt.imshow(im)
                plt.subplot(2, 2, 2)
                plt.imshow(np.clip(image_target_disp * 255, 0, 255).astype(np.uint8))
                plt.subplot(2, 2, 3)
                plt.imshow(depth_disp)
                plt.subplot(2, 2, 4)
                plt.imshow(depth_target_disp)
                plt.show()

            # AAE forward pass
            images_input = torch.cat((images, class_info), dim=1)
            outputs = self.AAE.forward(images_input)
            images_recnst = outputs[0]
            loss_reconstr = self.AAE.B_loss(images_recnst[:, :3, :, :], images_target.detach())
            loss = loss_reconstr

            loss_aae_rgb_data = loss_reconstr.data.cpu().item()
            loss_aae_depth_data = 0

            # AAE backward pass
            optimizer.zero_grad()
            try:
                loss.backward()
            except:
                pass

            optimizer.step()

            print("{}/{}: {}/{}, loss_rgb: {:.4f}, loss depth: {:.4f}".format(epoch, total_epoch, step + 1, int(steps),
                  loss_aae_rgb_data, loss_aae_depth_data))

            loss_sum_rgb += loss_aae_rgb_data / steps
            loss_sum_depth += loss_aae_depth_data / steps

            # display
            plot_n_comparison = 20
            if step < plot_n_comparison:
                image = images[0].detach().permute(1, 2, 0).cpu().numpy()
                image_target = images_target[0].permute(1, 2, 0).cpu().numpy()
                image_recon = images_recnst[0, :3].detach().permute(1, 2, 0).cpu().numpy()
                disp = (image, image_target, image_recon)
                self.plot_comparison(disp, str(step))

            if step==steps-1:
                break
            step += 1

        return loss_sum_rgb, loss_sum_depth

    # visualization
    def plot_comparison(self, images, name, save=True):
        if len(images) == 3:
            comparison_row = np.concatenate(images, 1)
            plt.figure("compare"+name)
            plt.gcf().clear()
            plt.imshow(comparison_row)
        else:
            comparison_row = np.concatenate(images[:3], 1)
            comparison_row2 = np.concatenate(images[3:], 1)
            plt.figure("compare" + name)
            plt.gcf().clear()
            plt.subplot(2, 1, 1)
            plt.imshow(comparison_row)
            plt.subplot(2, 1, 2)
            plt.imshow(comparison_row2)

        if save:
            save_path = os.path.join(self.log_dir, "compare_{}.png".format(name))
            plt.savefig(save_path)
            plt.close()
        else:
            plt.show(block=False)
            plt.pause(0.1)


import os
import logging
import tempfile
import subprocess

import tensorflow as tf
import numpy as np

import keras_facenet

logging.basicConfig()
log = logging.getLogger()
log.setLevel('INFO')

facenet_data_dir = 'tests/facenet_data_dir'
keras_facenet_data_dir = 'tests/keras_facenet_data_dir'
os.makedirs(facenet_data_dir, exist_ok=True)
os.makedirs(keras_facenet_data_dir, exist_ok=True)
images = np.random.randn(3, 160, 160, 3)

subprocess.call(['git', 'submodule', 'init'])
subprocess.call(['git', 'submodule', 'update'])

from .facenet.src import download_and_extract, facenet  # noqa: E402

def load_model(model):
    # This is a modified version of facenet.load_model to support TensorFlow 2.x.
    model_exp = os.path.expanduser(model)
    meta_file, ckpt_file = facenet.get_model_filenames(model_exp)
    saver = tf.compat.v1.train.import_meta_graph(os.path.join(model_exp, meta_file))
    saver.restore(tf.compat.v1.get_default_session(), os.path.join(model_exp, ckpt_file))

def run_test(model_name):
    download_and_extract.download_and_extract_file(
        model_name=model_name,
        data_dir=facenet_data_dir
    )
    kfn = keras_facenet.FaceNet(
        key=model_name,
        cache_folder=keras_facenet_data_dir
    )
    emb_kfn = kfn.embeddings(images)
    with tf.Graph().as_default():
        with tf.compat.v1.Session() as sess:
            load_model(os.path.join(facenet_data_dir, model_name))
            images_placeholder = tf.compat.v1.get_default_graph().get_tensor_by_name("input:0")  # noqa: E501
            embeddings = tf.compat.v1.get_default_graph().get_tensor_by_name("embeddings:0")  # noqa: E501
            phase_train_placeholder = tf.compat.v1.get_default_graph().get_tensor_by_name("phase_train:0")  # noqa: E501
            if kfn.metadata['fixed_image_standardization']:
                X = (images - 127.5) / 127.5
            else:
                X = np.array([facenet.prewhiten(image) for image in images])
            feed_dict = {images_placeholder: X, phase_train_placeholder: False}  # noqa: E501
            emb = sess.run(embeddings, feed_dict=feed_dict)
    print('Comparing results for', model_name)
    np.testing.assert_almost_equal(emb, emb_kfn)

def test_2017_weights():
    run_test('20170512-110547')

def test_2018_weights():
    run_test('20180402-114759')


import torch
import numpy as np
import time
import pdb
from rlkit.torch.core import eval_np
from rlkit.torch.policies.tanh_gaussian_policy import TanhGaussianPolicy
from rlkit.torch.policies.make_deterministic import MakeDeterministic
from rlkit.torch.policies.gumbel_softmax_policy import GumbelSoftmaxMlpPolicy
from rlkit.policies.argmax import ArgmaxDiscretePolicy

import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--exp_name', type=str, default='simple_spread')
parser.add_argument('--boundary', action='store_true', default=False)
parser.add_argument('--sb', action='store_true', default=False) # shared obs
parser.add_argument('--num_ag', type=int, default=None)
parser.add_argument('--num_adv', type=int, default=None)
parser.add_argument('--num_l', type=int, default=None)
parser.add_argument('--ar', type=float, default=None) # agent radius
parser.add_argument('--sr', action='store_true', default=False) # shape reward
parser.add_argument('--sg', action='store_true', default=False) # shared groups
parser.add_argument('--mpl', type=int, default=25) # max path length
parser.add_argument('--log_dir', type=str, default='MASAC')
parser.add_argument('--epoch', type=int, default=None)
parser.add_argument('--seed', type=str, default=0)
args = parser.parse_args()

pre_dir = './Data/'+args.exp_name\
            +('bd' if args.boundary else '')\
            +('sr' if args.sr else '')\
            +(('ag'+str(args.num_ag)) if args.num_ag else '')\
            +(('adv'+str(args.num_adv)) if args.num_adv else '')\
            +(('l'+str(args.num_l)) if args.num_l else '')\
            +(('ar'+str(args.ar)) if args.ar else '')\
            +'_mpl'+str(args.mpl)
if args.epoch:
    data_path = '{}/{}/seed{}/itr_{}.pkl'.format(pre_dir,args.log_dir,args.seed,args.epoch)
else:
    data_path = '{}/{}/seed{}/params.pkl'.format(pre_dir,args.log_dir,args.seed)
data = torch.load(data_path,map_location='cpu')
policy_n = data['trainer/policy_n']
if isinstance(policy_n[0],TanhGaussianPolicy):
    policy_n = [MakeDeterministic(policy) for policy in policy_n]
elif  isinstance(policy_n[0],GumbelSoftmaxMlpPolicy):
    policy_n = [ArgmaxDiscretePolicy(policy,use_preactivation=True) for policy in policy_n]

if 'trainer/shared_gnn' in data.keys():
    shared_gnn = data['trainer/shared_gnn']
elif 'R2GGNN12Share' in args.log_dir:
    shared_gnn = data['trainer/qf1'].obs_gnn
else:
    shared_gnn = None

world_args=dict(
    num_agents=args.num_ag,
    num_adversaries=args.num_adv,
    num_landmarks=args.num_l,
    agent_size=args.ar,
    boundary=([[-1.,-1.],[1.,1.]] if args.boundary else None),
    shape_reward=args.sr,
    shared_obs=args.sb,
)
import sys
sys.path.append("./particle-graph-envs")
from make_env import make_env
env = make_env(args.exp_name,discrete_action_space=False,world_args=world_args)
num_agent = env.num_agents
obs_dim = env.observation_space.low.size
action_dim = env.action_space.low.size

if args.sg:
    from build_networks import build_shared_groups
    shared_groups = build_shared_groups(args)
else:
    shared_groups = dict()
    for i in range(num_agent):
        shared_groups[i] = 0
print('shared_groups: ', shared_groups)

max_path_length = args.mpl
path_length = 0
done = np.array([False]*num_agent)
c_r = np.zeros(num_agent)
o_n = env.reset()
while True:
    path_length += 1
    a_n = []
    if args.sb:
        if shared_gnn:
            if args.sg:
                o_emb_n = [eval_np(sub_shared_encoder,o_n[None,:])[0] for sub_shared_encoder in shared_gnn]
                for i, policy in enumerate(policy_n):
                    o = o_emb_n[shared_groups[i]][i]
                    a, info = policy.get_action(o)
                    a_n.append(a)
            else:
                if isinstance(shared_gnn, list):
                    assert len(shared_gnn) == 1
                    shared_gnn = shared_gnn[0]
                o_emb_n = eval_np(shared_gnn,o_n[None,:])[0]
                for i,policy in enumerate(policy_n):
                    o = o_emb_n[i]
                    a, info = policy.get_action(o)
                    a_n.append(a)
        else:
            for policy in policy_n:
                a, info = policy.get_action(o_n)
                a_n.append(a)
    else:
        for o,policy in zip(o_n,policy_n):
            a, info = policy.get_action(o)
            a_n.append(a)

    o_n, r_n, done, _ = env.step(a_n)
    c_r += r_n
    env.render()
    print("step: ",path_length)
    print("a: ",a_n)
    print("info: ",info)
    print("o: ",o_n)
    print('r: ',r_n)
    print(done)
    time.sleep(0.1)
    if path_length > max_path_length or done.all():
        print('c_r: ',c_r)
        path_length = 0
        done = np.array([False]*num_agent)
        c_r = np.zeros(num_agent)
        o_n = env.reset()
        pdb.set_trace()


import sys
import warnings

import numpy as np
import pandas as pd
import torch

import datasets
import models
import training
import utils

from pathlib import Path

from tqdm import tqdm

from callbacks import lwlrap

if __name__ == "__main__":
    warnings.filterwarnings("ignore")

    ##################################################
    #  Basic configuration #
    ##################################################
    args = utils.get_parser().parse_args()
    config = utils.load_config(args.config)

    global_params = config["globals"]

    # logging
    config_name = args.config.split("/")[-1].replace(".yml", "")
    expdir = Path(f"out/{config_name}")
    if not expdir.exists():
        print(f"You need to train {config_name} first!")
        sys.exit(1)
    submission_file_dir = expdir / "submission"
    submission_file_dir.mkdir(parents=True, exist_ok=True)

    logger = utils.get_logger(expdir / "tta.log")

    # environment
    utils.set_seed(global_params["seed"])
    device = training.get_device(global_params["device"])

    # data
    tp, fp, train_all, test_all, train_audio, test_audio = datasets.get_metadata(config)
    submission = pd.read_csv(config["data"]["sample_submission_path"])

    # labels
    labels = []
    duration = config["dataset"]["valid"]["params"]["duration"]
    for _, sample in tp.iterrows():
        t_min = sample["t_min"]
        t_max = sample["t_max"]
        flac_id = sample["recording_id"]
        call_duration = t_max - t_min
        relative_offset = (duration - call_duration) / 2

        offset = min(max(0, t_min - relative_offset), 60 - duration)
        tail = offset + duration

        query_string = f"recording_id == '{flac_id}' & "
        query_string += f"t_min < {tail} & t_max > {offset}"
        all_tp_events = tp.query(query_string)

        label = np.zeros(24, dtype=np.float32)
        for species_id in all_tp_events["species_id"].unique():
            label[int(species_id)] = 1.0
        labels.append(label)

    labels_df = pd.DataFrame(np.asarray(labels), columns=[f"s{i}" for i in range(24)])
    ground_truth_df = pd.concat([
        pd.DataFrame({"index": tp["index"], "recording_id": tp["recording_id"]}),
        labels_df
    ], axis=1)

    # validation
    splitter = training.get_split(config)

    ##################################################
    # Main Loop #
    ##################################################
    fold_predictions = []
    oof_predictions = []
    oof_tta_dict: dict = {}
    for tta_conf in config["tta"]:
        oof_tta_dict[tta_conf["name"]] = []
    for i, (trn_idx, val_idx) in enumerate(splitter.split(train_all)):
        if i not in global_params["folds"]:
            continue
        logger.info("=" * 20)
        logger.info(f"Fold {i}")
        logger.info("=" * 20)

        val_df = train_all.iloc[val_idx, :].reset_index(drop=True)
        model = models.get_model(config)
        if config["inference"].get("last", False):
            model = models.prepare_for_inference(
                model, expdir / f"fold{i}/checkpoints/last.pth").to(device)
            last = True
        else:
            model = models.prepare_for_inference(
                model, expdir / f"fold{i}/checkpoints/best.pth").to(device)
            last = False

        ttas = config["tta"]
        oof_tta_predictions = []
        tta_predictions = []
        for tta in ttas:
            logger.info("#" * 20)
            logger.info(tta["name"])

            _config = config.copy()
            _config["transforms"]["valid"] = [tta]
            val_loader = datasets.get_train_loader(
                val_df, tp, fp, train_audio, _config, phase="valid")

            _config["transforms"]["test"] = [tta]

            loader = datasets.get_test_loader(test_all, test_audio, _config)

            if config["inference"]["prediction_type"] == "strong":
                ##################################################
                # OOF #
                ##################################################
                logger.info("*" * 20)
                logger.info(f"OOF prediction for fold{i}")
                logger.info("*" * 20)
                recording_ids = []
                batch_predictions = []
                indices = []
                for batch in tqdm(val_loader, leave=True):
                    recording_ids.extend(batch["recording_id"])
                    indices.extend(batch["index"].numpy())
                    input_ = batch[global_params["input_key"]].to(device)
                    with torch.no_grad():
                        output = model(input_)
                    framewise_output = output["framewise_output"].detach()
                    clipwise_output, _ = framewise_output.max(dim=1)
                    batch_predictions.append(
                        clipwise_output.cpu().numpy())
                oof_prediction = np.concatenate(batch_predictions, axis=0)
                oof_prediction_df = pd.DataFrame(
                    oof_prediction, columns=[f"s{i}" for i in range(oof_prediction.shape[1])])
                oof_prediction_df = pd.concat([
                    pd.DataFrame({"index": indices, "recording_id": recording_ids}),
                    oof_prediction_df
                ], axis=1)

                oof_tta_predictions.append(oof_prediction_df)
                oof_tta_dict[tta["name"]].append(oof_prediction_df)
                if last:
                    oof_name = "oof_tta_strong_last.csv"
                else:
                    oof_name = "oof_tta_strong.csv"

                ##################################################
                # Prediction #
                ##################################################
                logger.info("*" * 20)
                logger.info(f"Prediction on test for fold{i}")
                logger.info("*" * 20)
                recording_ids = []
                batch_predictions = []
                for batch in tqdm(loader, leave=True):
                    recording_ids.extend(batch["recording_id"])
                    input_ = batch[global_params["input_key"]].to(device)
                    with torch.no_grad():
                        output = model(input_)
                    framewise_output = output["framewise_output"].detach()
                    clipwise_output, _ = framewise_output.max(dim=1)
                    batch_predictions.append(
                        clipwise_output.cpu().numpy())
                fold_prediction = np.concatenate(batch_predictions, axis=0)

                fold_prediction_df = pd.DataFrame(
                    fold_prediction, columns=[f"s{i}" for i in range(fold_prediction.shape[1])])
                fold_prediction_df = pd.concat([
                    pd.DataFrame({"recording_id": recording_ids}),
                    fold_prediction_df
                ], axis=1)

                fold_prediction_df = fold_prediction_df.groupby(
                    "recording_id").max().reset_index(drop=False)
                tta_predictions.append(fold_prediction_df)
                if last:
                    submission_name = "tta_strong_last.csv"
                else:
                    submission_name = "tta_strong.csv"
            else:
                ##################################################
                # OOF #
                ##################################################
                logger.info("*" * 20)
                logger.info(f"OOF prediction for fold{i}")
                logger.info("*" * 20)
                recording_ids = []
                batch_predictions = []
                indices = []
                for batch in tqdm(val_loader, leave=True):
                    recording_ids.extend(batch["recording_id"])
                    indices.extend(batch["index"].numpy())
                    input_ = batch[global_params["input_key"]].to(device)
                    with torch.no_grad():
                        output = model(input_)
                    batch_predictions.append(
                        output["clipwise_output"].cpu().numpy())
                oof_prediction = np.concatenate(batch_predictions, axis=0)
                oof_prediction_df = pd.DataFrame(
                    oof_prediction, columns=[f"s{i}" for i in range(oof_prediction.shape[1])])
                oof_prediction_df = pd.concat([
                    pd.DataFrame({"index": indices, "recording_id": recording_ids}),
                    oof_prediction_df
                ], axis=1)

                oof_tta_predictions.append(oof_prediction_df)
                oof_tta_dict[tta["name"]].append(oof_prediction_df)
                if last:
                    oof_name = "oof_tta_weak_last.csv"
                else:
                    oof_name = "oof_tta_weak.csv"

                ##################################################
                # Prediction #
                ##################################################
                logger.info("*" * 20)
                logger.info(f"Prediction on test for fold{i}")
                logger.info("*" * 20)
                recording_ids = []
                batch_predictions = []
                for batch in tqdm(loader, leave=True):
                    recording_ids.extend(batch["recording_id"])
                    input_ = batch[global_params["input_key"]].to(device)
                    with torch.no_grad():
                        output = model(input_)
                    batch_predictions.append(
                        output["clipwise_output"].detach().cpu().numpy())
                fold_prediction = np.concatenate(batch_predictions, axis=0)

                fold_prediction_df = pd.DataFrame(
                    fold_prediction, columns=[f"s{i}" for i in range(fold_prediction.shape[1])])
                fold_prediction_df = pd.concat([
                    pd.DataFrame({"recording_id": recording_ids}),
                    fold_prediction_df
                ], axis=1)

                fold_prediction_df = fold_prediction_df.groupby(
                    "recording_id").max().reset_index(drop=False)
                tta_predictions.append(fold_prediction_df)
                if last:
                    submission_name = "tta_weak_last.csv"
                else:
                    submission_name = "tta_weak.csv"
        oof_tta_df = pd.concat(oof_tta_predictions, axis=0).reset_index(drop=True)
        oof_tta_df = oof_tta_df.groupby(["index", "recording_id"])[[f"s{i}" for i in range(24)]].mean().reset_index(drop=False)
        oof_predictions.append(oof_tta_df)

        tta_pred_df = pd.concat(tta_predictions, axis=0).reset_index(drop=True)
        tta_pred_df = tta_pred_df.groupby(["recording_id"])[[f"s{i}" for i in range(24)]].mean().reset_index(drop=False)
        fold_predictions.append(tta_pred_df)

    oof_df = pd.concat(oof_predictions, axis=0).reset_index(drop=True)

    oof_indices = oof_df[["index"]]
    ground_truth = oof_indices.merge(ground_truth_df, on="index", how="left")
    columns = [f"s{i}" for i in range(24)]
    score_class, weight = lwlrap(ground_truth[columns].values, oof_df[columns].values)
    score = (score_class * weight).sum()
    logger.info(f"TTA all LWLRAP: {score:.5f}")

    class_level_score = {config_name: score_class}

    for key in oof_tta_dict:
        tta_df = pd.concat(oof_tta_dict[key], axis=0).reset_index(drop=True)
        oof_indices = tta_df[["index"]]
        ground_truth = oof_indices.merge(ground_truth_df, on="index", how="left")
        score_class, weight = lwlrap(ground_truth[columns].values, tta_df[columns].values)
        score = (score_class * weight).sum()
        logger.info(f"TTA {key} LWLRAP: {score:.5f}")
        class_level_score[key] = score_class

    utils.save_json(class_level_score, submission_file_dir / "class_level_results.json")

    oof_df.to_csv(submission_file_dir / oof_name, index=False)

    folds_prediction_df = pd.concat(fold_predictions, axis=0).reset_index(drop=True)
    folds_prediction_df = folds_prediction_df.groupby("recording_id").mean().reset_index(drop=False)

    assert len(folds_prediction_df) == len(submission), \
        "prediction length does not match sample submission length"
    assert folds_prediction_df.shape[1] == submission.shape[1], \
        "number of classes in prediction does not match that of sample submission"
    assert len(set(folds_prediction_df["recording_id"]) - set(submission["recording_id"])) == 0, \
        "recording_id in prediction has unknown value"
    assert len(set(submission["recording_id"]) - set(folds_prediction_df["recording_id"])) == 0, \
        "prediction doesn't have enough recording_id"

    folds_prediction_df.to_csv(submission_file_dir / submission_name, index=False)


import timm
import numpy as np
import torch.nn as nn

from ._base import EncoderMixin

def _make_divisible(x, divisible_by=8):
    return int(np.ceil(x * 1.0 / divisible_by) * divisible_by)

class MobileNetV3Encoder(nn.Module, EncoderMixin):
    def __init__(self, model_name, width_mult, depth=5, **kwargs):
        super().__init__()
        if "large" not in model_name and "small" not in model_name:
            raise ValueError("MobileNetV3 wrong model name {}".format(model_name))

        self._mode = "small" if "small" in model_name else "large"
        self._depth = depth
        self._out_channels = self._get_channels(self._mode, width_mult)
        self._in_channels = 3

        # minimal models replace hardswish with relu
        self.model = timm.create_model(
            model_name=model_name,
            scriptable=True,  # torch.jit scriptable
            exportable=True,  # onnx export
            features_only=True,
        )

    def _get_channels(self, mode, width_mult):
        if mode == "small":
            channels = [16, 16, 24, 48, 576]
        else:
            channels = [16, 24, 40, 112, 960]
        channels = [
            3,
        ] + [_make_divisible(x * width_mult) for x in channels]
        return tuple(channels)

    def get_stages(self):
        if self._mode == "small":
            return [
                nn.Identity(),
                nn.Sequential(
                    self.model.conv_stem,
                    self.model.bn1,
                    self.model.act1,
                ),
                self.model.blocks[0],
                self.model.blocks[1],
                self.model.blocks[2:4],
                self.model.blocks[4:],
            ]
        elif self._mode == "large":
            return [
                nn.Identity(),
                nn.Sequential(
                    self.model.conv_stem,
                    self.model.bn1,
                    self.model.act1,
                    self.model.blocks[0],
                ),
                self.model.blocks[1],
                self.model.blocks[2],
                self.model.blocks[3:5],
                self.model.blocks[5:],
            ]
        else:
            ValueError("MobileNetV3 mode should be small or large, got {}".format(self._mode))

    def forward(self, x):
        stages = self.get_stages()

        features = []
        for i in range(self._depth + 1):
            x = stages[i](x)
            features.append(x)

        return features

    def load_state_dict(self, state_dict, **kwargs):
        state_dict.pop("conv_head.weight", None)
        state_dict.pop("conv_head.bias", None)
        state_dict.pop("classifier.weight", None)
        state_dict.pop("classifier.bias", None)
        self.model.load_state_dict(state_dict, **kwargs)

mobilenetv3_weights = {
    "tf_mobilenetv3_large_075": {
        "imagenet": "https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_large_075-150ee8b0.pth"  # noqa
    },
    "tf_mobilenetv3_large_100": {
        "imagenet": "https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_large_100-427764d5.pth"  # noqa
    },
    "tf_mobilenetv3_large_minimal_100": {
        "imagenet": "https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_large_minimal_100-8596ae28.pth"  # noqa
    },
    "tf_mobilenetv3_small_075": {
        "imagenet": "https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_small_075-da427f52.pth"  # noqa
    },
    "tf_mobilenetv3_small_100": {
        "imagenet": "https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_small_100-37f49e2b.pth"  # noqa
    },
    "tf_mobilenetv3_small_minimal_100": {
        "imagenet": "https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_small_minimal_100-922a7843.pth"  # noqa
    },
}

pretrained_settings = {}
for model_name, sources in mobilenetv3_weights.items():
    pretrained_settings[model_name] = {}
    for source_name, source_url in sources.items():
        pretrained_settings[model_name][source_name] = {
            "url": source_url,
            "input_range": [0, 1],
            "mean": [0.485, 0.456, 0.406],
            "std": [0.229, 0.224, 0.225],
            "input_space": "RGB",
        }

timm_mobilenetv3_encoders = {
    "timm-mobilenetv3_large_075": {
        "encoder": MobileNetV3Encoder,
        "pretrained_settings": pretrained_settings["tf_mobilenetv3_large_075"],
        "params": {"model_name": "tf_mobilenetv3_large_075", "width_mult": 0.75},
    },
    "timm-mobilenetv3_large_100": {
        "encoder": MobileNetV3Encoder,
        "pretrained_settings": pretrained_settings["tf_mobilenetv3_large_100"],
        "params": {"model_name": "tf_mobilenetv3_large_100", "width_mult": 1.0},
    },
    "timm-mobilenetv3_large_minimal_100": {
        "encoder": MobileNetV3Encoder,
        "pretrained_settings": pretrained_settings["tf_mobilenetv3_large_minimal_100"],
        "params": {"model_name": "tf_mobilenetv3_large_minimal_100", "width_mult": 1.0},
    },
    "timm-mobilenetv3_small_075": {
        "encoder": MobileNetV3Encoder,
        "pretrained_settings": pretrained_settings["tf_mobilenetv3_small_075"],
        "params": {"model_name": "tf_mobilenetv3_small_075", "width_mult": 0.75},
    },
    "timm-mobilenetv3_small_100": {
        "encoder": MobileNetV3Encoder,
        "pretrained_settings": pretrained_settings["tf_mobilenetv3_small_100"],
        "params": {"model_name": "tf_mobilenetv3_small_100", "width_mult": 1.0},
    },
    "timm-mobilenetv3_small_minimal_100": {
        "encoder": MobileNetV3Encoder,
        "pretrained_settings": pretrained_settings["tf_mobilenetv3_small_minimal_100"],
        "params": {"model_name": "tf_mobilenetv3_small_minimal_100", "width_mult": 1.0},
    },
}


"""

Test scripts

"""
import argparse
import json
import tqdm
import torch
import os
import numpy as np
from lib.det_ops.anchors import compute_anchor_whs, generate_anchors
from lib.bbox import bbox, box_transform
from lib.nms import nms
from utils.logger import load_checkpoints
from models import retina
from IPython import embed
from datasets import voc, coco, images
from cfgs import config as cfg
from torchvision.datasets import ImageFolder

import warnings

warnings.filterwarnings("ignore", category=DeprecationWarning)

def inference(model, dataset, anchor_wh, strides, result_file, config):

    model.eval()
    num_samples = len(dataset)
    pbar = tqdm.tqdm(range(num_samples))

    if os.path.exists(result_file):
        try:
            os.remove(result_file)
        except:
            print("Error while deleting current file and opening new file: ", result_file)

    results = []

    with torch.no_grad():
        for i, idx in enumerate(pbar):
            img, im_name, scale, im_size = dataset[idx]

            img = img.cuda()
            cls_pred, bbox_pred = model(img.unsqueeze(0))
            scores = cls_pred.sigmoid()
            # bbox [N, 4]
            bbox_pred = bbox_pred[0]
            # cls [N, C]
            scores = scores[0]

            h, w = img.shape[1], img.shape[2]

            anchors = generate_anchors(anchor_wh, input_size=np.array([h, w]),
                                       strides=strides)
            anchors = anchors.cuda()

            # transform to bboxes
            boxes = box_transform.bbox_transform_inv(anchors, bbox_pred.squeeze())
            boxes = boxes/scale

            filter_boxes_inds_x = boxes[:, 0] >= boxes[:, 2]
            filter_boxes_inds_y = boxes[:, 1] >= boxes[:, 3]
            filter_boxes_inds = (1 - filter_boxes_inds_x) * (1 - filter_boxes_inds_y)
            boxes = boxes[filter_boxes_inds]
            scores = scores[filter_boxes_inds]

            scores = scores.squeeze()

            result_boxes = []  # []
            # every class
            # 1. max detection score
            # 2. score thresh
            # 3. do nms
            # 4. top k
            max_labels = torch.argmax(scores, dim=1)

            for cls in range(config['num_classes']-1):

                # filter predictions through 'classification threshold'
                score = scores[:, cls]
                cls_inds = score > config['cls_thresh']
                # current class has the max score over all classes
                max_inds = max_labels == cls
                cls_inds = max_inds * cls_inds
                if cls_inds.sum() < 1:
                    continue
                # score [K]
                score = score[cls_inds]

                # _boxes [K, 4]
                _boxes = boxes[cls_inds]

                # NMS remove duplicate
                keep = nms(torch.cat([_boxes, score.unsqueeze(1)], 1), config['test_nms'])

                score = score[keep]
                _boxes = _boxes[keep]

                for i in range(_boxes.shape[0]):
                    result_boxes.append((cls, score[i].item(), _boxes[i].cpu().data.numpy().tolist()))

            # Keep Max Num Boxes
            if len(result_boxes) > config['test_max_boxes']:
                result_boxes = sorted(result_boxes, key=lambda x: x[1], reverse=True)
                result_boxes = result_boxes[:config['test_max_boxes']]
            pbar.set_description('im_det:{}/{}'.format(idx, num_samples))

            if len(result_boxes) == 0:
                continue

            for j in range(len(result_boxes)):
                cls, s, b, = result_boxes[j]
                if s > 0.75:
                    pred = dict()

                    pred['image_id'] = im_name
                    pred['score'] = s
                    pred['category_id'] = cls + 1
                    pred['bbox'] = b
                    results.append(pred)

    with open(result_file, 'a+') as f:
        s = json.dumps(results)
        f.write(s)
    print("JSON Finished!")
    print('File written to: ', result_file)

def load_img(path):
    npi = imread(path)
    return Image.fromarray(npi)

def validate(args, config):

    anchor_scales = config['anchor_sizes']
    anchor_apsect_ratios = config['aspect_ratios']
    num_anchors = len(anchor_scales) * len(anchor_apsect_ratios)

    model = retina.RetinaNet(config['num_classes']-1, num_anchors, config['basemodel_path']).cuda()

    model_path = args.model_path
    print('Model Path: ', model_path)
    output_file = args.output
    if args.dataset == 'VOC':
        dataset = voc.VOC2012(dataroot=config['data_dir'], imageset=args.imageset, config=config)
    elif args.dataset == 'COCO':
        dataset = coco.COCODetection(dataroot=config['data_dir'], imageset=args.imageset, config=config, training=False)
    elif args.dataset == 'images':
        dataset = images.ImageData(dataroot=args.path, config=config)
    else:
        raise NotImplemented()
    print('Dataset: ', len(dataset))
    state_dict, _, _, _ = load_checkpoints(model_path)
    model.load_state_dict(state_dict)

    anchor_whs = compute_anchor_whs(len(config['strides']), areas=config['anchor_areas'],
                                    aspect_ratios=anchor_apsect_ratios,
                                    sizes=anchor_scales)

    inference(model, dataset, anchor_whs, config['strides'], result_file=output_file, config=config)

if __name__ == '__main__':
    os.environ["CUDA_VISIBLE_DEVICES"] = '2'

    parser = argparse.ArgumentParser('Test RetinaNet')
    parser.add_argument('-o', '--output', type=str, default='./result.json', help='output file path')
    parser.add_argument('-m', '--model_path', type=str, help='saved model path')
    parser.add_argument('-i', '--imageset', type=str, default='val', help='imageset to predict')
    parser.add_argument('-p', '--path', type=str, default='',
                        help='path, where to find the image set')
    parser.add_argument('-e', '--experiment', type=str, default='coco_baseline',
                        help='experiment name, correspond to `config.py`')
    parser.add_argument('-ds', '--dataset', type=str, default='COCO', help='dataset')

    _args = parser.parse_args()
    config = cfg.config[_args.experiment]
    _args = parser.parse_args()
    validate(_args, config)


import os
import torch
import numpy as np
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
import config as cfg
import random
import pickle
import parameters_BL as params
import json
import math
import cv2
from tqdm import tqdm
import time
import torchvision.transforms as trans
from decord import VideoReader

class baseline_dataloader_train_strong(Dataset):

    def __init__(self, shuffle = True, data_percentage = 1.0, split = 1):

        if split == 1:
            self.all_paths = open(os.path.join(cfg.path_folder, 'ucfTrainTestlist/trainlist01.txt'),'r').read().splitlines()
        elif split ==2:
            self.all_paths = open(os.path.join(cfg.path_folder, 'ucfTrainTestlist/trainlist02.txt'),'r').read().splitlines()
        elif split ==3:
            self.all_paths = open(os.path.join(cfg.path_folder, 'ucfTrainTestlist/trainlist03.txt'),'r').read().splitlines()
        else:
            print(f'Invalid split input: {split}')

        self.classes= json.load(open(cfg.class_mapping))['classes']
        self.shuffle = shuffle

        if self.shuffle:
            random.shuffle(self.all_paths)

        self.data_percentage = data_percentage
        self.data_limit = int(len(self.all_paths)*self.data_percentage)
        self.data = self.all_paths[0: self.data_limit]
        self.PIL = trans.ToPILImage()
        self.TENSOR = trans.ToTensor()
        self.erase_size = 19

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

    def __getitem__(self,index):
        clip, label, vid_path = self.process_data(index)
        return clip, label, vid_path

    def process_data(self, idx):

        # label_building
        vid_path = cfg.path_folder + '/UCF-101/' + self.data[idx].split(' ')[0]

        label = self.classes[vid_path.split('/')[-2]] # This element should be activity name

        # clip_building
        clip = self.build_clip(vid_path)

        return clip, label, vid_path
    def build_clip(self, vid_path):

        try:
            cap = cv2.VideoCapture(vid_path)
            cap.set(1, 0)
            frame_count = cap.get(7)

            ############################# frame_list maker start here#################################

            skip_frames_full = params.fix_skip #frame_count/(params.num_frames)

            left_over = frame_count - params.fix_skip*params.num_frames

            start_frame_full = np.random.randint(0,int(left_over))

            frames_full = start_frame_full + np.asarray([int(int(skip_frames_full)*f) for f in range(params.num_frames)])

            if frames_full[-1] >= frame_count:
                # print('some corner case not covered')
                frames_full[-1] = int(frame_count-1)
            ################################ frame list maker finishes here ###########################

            ################################ actual clip builder starts here ##########################
            full_clip = []
            list_full = []
            count = -1
            random_array = np.random.rand(2,8)
            x_erase = np.random.randint(0,params.reso_h, size = (2,))
            y_erase = np.random.randint(0,params.reso_w, size = (2,))

            cropping_factor1 = np.random.uniform(0.6, 1, size = (2,)) # on an average cropping factor is 80% i.e. covers 64% area
            x0 = np.random.randint(0, params.ori_reso_w - params.ori_reso_w*cropping_factor1[0] + 1)
            y0 = np.random.randint(0, params.ori_reso_h - params.ori_reso_h*cropping_factor1[0] + 1)

            #Here augmentations are not strong as self-supervised training
            contrast_factor1 = np.random.uniform(0.9,1.1, size = (2,))
            hue_factor1 = np.random.uniform(-0.05,0.05, size = (2,))
            saturation_factor1 = np.random.uniform(0.9,1.1, size = (2,))
            brightness_factor1 = np.random.uniform(0.9,1.1,size = (2,))
            gamma1 = np.random.uniform(0.85,1.15, size = (2,))

            erase_size1 = np.random.randint(int(self.erase_size/2),self.erase_size, size = (2,))
            erase_size2 = np.random.randint(int(self.erase_size/2),self.erase_size, size = (2,))
            random_color_dropped = np.random.randint(0,3,(2))

            while(cap.isOpened()):
                count += 1
                ret, frame = cap.read()
                if ((count not in frames_full)) and (ret == True):
                    continue
                if ret == True:
                    if (count in frames_full):
                        full_clip.append(self.augmentation(frame, random_array[0], x_erase[0], y_erase[0], cropping_factor1[0],\
                                x0, y0, contrast_factor1[0], hue_factor1[0], saturation_factor1[0], brightness_factor1[0],\
                                gamma1[0],erase_size1[0],erase_size2[0], random_color_dropped[0]))
                        list_full.append(count)
                else:
                    break

            if len(full_clip) < params.num_frames and len(full_clip)>(params.num_frames/2) :
                print(f'Clip {vid_path} is missing {params.num_frames-len(full_clip)} frames')
                remaining_num_frames = params.num_frames - len(full_clip)
                full_clip = full_clip + full_clip[::-1][1:remaining_num_frames+1]

            try:
                assert(len(full_clip)==params.num_frames)

                return full_clip
            except:
                print(frames_full)
                print(f'Clip {vid_path} Failed')
                return None

        except:
            print(f'Clip {vid_path} Failed')
            return None

    def augmentation(self, image, random_array, x_erase, y_erase, cropping_factor1,\
        x0, y0, contrast_factor1, hue_factor1, saturation_factor1, brightness_factor1,\
        gamma1,erase_size1,erase_size2, random_color_dropped):

        image = self.PIL(image)
        image = trans.functional.resized_crop(image,y0,x0,int(params.ori_reso_h*cropping_factor1),int(params.ori_reso_h*cropping_factor1),(params.reso_h,params.reso_w),interpolation=2)

        if random_array[0] < 0.125/2:
            image = trans.functional.adjust_contrast(image, contrast_factor = contrast_factor1) #0.75 to 1.25
        if random_array[1] < 0.3/2 :
            image = trans.functional.adjust_hue(image, hue_factor = hue_factor1)
        if random_array[2] < 0.3/2 :
            image = trans.functional.adjust_saturation(image, saturation_factor = saturation_factor1)
        if random_array[3] < 0.3/2 :
            image = trans.functional.adjust_brightness(image, brightness_factor = brightness_factor1)
        if random_array[0] > 0.125/2 and random_array[0] < 0.25/2:
            image = trans.functional.adjust_contrast(image, contrast_factor = contrast_factor1) #0.75 to 1.25
        if random_array[4] > 0.9:
            image = trans.functional.to_grayscale(image, num_output_channels = 3)
            if random_array[5] > 0.25:
                image = trans.functional.adjust_gamma(image, gamma = gamma1, gain=1)
        if random_array[6] > 0.5:
            image = trans.functional.hflip(image)

        image = trans.functional.to_tensor(image)

        if random_array[7] < 0.5/2 :
            image = trans.functional.erase(image, x_erase, y_erase, erase_size1, erase_size2, v=0)

        return image

class multi_baseline_dataloader_val_strong(Dataset):

    def __init__(self, shuffle = True, data_percentage = 1.0, mode = 0, \
                hflip=0, cropping_factor=0.8):
        if split == 1:
            self.all_paths = open(os.path.join(cfg.path_folder, 'ucfTrainTestlist/testlist01.txt'),'r').read().splitlines()
        elif split ==2:
            self.all_paths = open(os.path.join(cfg.path_folder, 'ucfTrainTestlist/testlist02.txt'),'r').read().splitlines()
        elif split ==3:
            self.all_paths = open(os.path.join(cfg.path_folder, 'ucfTrainTestlist/testlist03.txt'),'r').read().splitlines()
        else:
            print(f'Invalid split input: {split}')

        self.classes= json.load(open(cfg.class_mapping))['classes']
        self.shuffle = shuffle

        if self.shuffle:
            random.shuffle(self.all_paths)

        self.data_percentage = data_percentage
        self.data_limit = int(len(self.all_paths)*self.data_percentage)
        self.data = self.all_paths[0: self.data_limit]
        self.PIL = trans.ToPILImage()
        self.TENSOR = trans.ToTensor()
        self.mode = mode
        self.hflip = hflip
        self.cropping_factor = cropping_factor

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

    def __getitem__(self,index):
        clip, label, vid_path, frame_list = self.process_data(index)
        return clip, label, vid_path, frame_list

    def process_data(self, idx):

        # label_building
        vid_path = cfg.path_folder + '/UCF-101/' + self.data[idx].split(' ')[0]
        label = self.classes[vid_path.split('/')[-2]] # THIS MIGHT BE DIFFERNT AFTER STEVE MOVE THE PATHS

        # clip_building
        clip, frame_list = self.build_clip(vid_path)

        return clip, label, vid_path, frame_list

    def build_clip(self, vid_path):
        try:
            cap = cv2.VideoCapture(vid_path)
            cap.set(1, 0)
            frame_count = cap.get(7)

            N = frame_count
            n = params.num_frames

            skip_frames_full = params.fix_skip

            left_over = skip_frames_full*n
            F = N - left_over

            start_frame_full = 0 + int(np.linspace(0,F-10,params.num_modes)[self.mode])

            if start_frame_full< 0:
                start_frame_full = self.mode

            full_clip_frames = []

            full_clip_frames = start_frame_full + np.asarray(
                [int(int(skip_frames_full) * f) for f in range(params.num_frames)])

            count = -1
            full_clip = []
            list_full = []

            while (cap.isOpened()):
                count += 1
                ret, frame = cap.read()
                if ((count not in full_clip_frames) and (ret == True)):
                    continue
                if ret == True:
                    if (count in full_clip_frames):
                        full_clip.append(self.augmentation(frame))
                        list_full.append(count)

                else:
                    break
            # Appending the remaining frames in case of clip length < required frames
            if len(full_clip) < params.num_frames and len(full_clip)>(params.num_frames/2):
                remaining_num_frames = params.num_frames - len(full_clip)
                full_clip = full_clip + full_clip[::-1][1:remaining_num_frames+1]
            assert (len(full_clip) == params.num_frames)

            return full_clip, list_full

        except:
            print(f'Clip {vid_path} Failed')
            return None, None

    def augmentation(self, image):
        image = self.PIL(image)

        if self.cropping_factor <= 1:
            image = trans.functional.center_crop(image,(int(params.ori_reso_h*self.cropping_factor),int(params.ori_reso_h*self.cropping_factor)))
        image = trans.functional.resize(image, (params.reso_h, params.reso_w))
        if self.hflip !=0:
            image = trans.functional.hflip(image)

        return trans.functional.to_tensor(image)

def collate_fn1(batch):
    clip, label, vid_path = [], [], []
    for item in batch:
        if not (item[0] == None or item[1] == None or item[2] == None):
            clip.append(torch.stack(item[0],dim=0))

            label.append(item[1])
            vid_path.append(item[2])

    clip = torch.stack(clip, dim=0)

    return clip, label, vid_path

def collate_fn2(batch):

    f_clip, label, vid_path, frame_list = [], [], [], []
    # print(len(batch))
    for item in batch:
        if not (item[0] == None or item[1] == None or item[2] == None):
            f_clip.append(torch.stack(item[0],dim=0))
            label.append(item[1])
            vid_path.append(item[2])
            frame_list.append(item[3])
    f_clip = torch.stack(f_clip, dim=0)

    return f_clip, label, vid_path, frame_list

def collate_fn_train(batch):

    f_clip, label, vid_path = [], [], [],
    for item in batch:
        if not (item[0] == None or item[1] == None or item[2] == None):
            f_clip.append(torch.stack(item[0],dim=0))
            label.append(item[1])
            vid_path.append(item[2])

    f_clip = torch.stack(f_clip, dim=0)

    return f_clip, label, vid_path

if __name__ == '__main__':

    train_dataset = baseline_dataloader_train_strong(shuffle = False, data_percentage = 1.0)
    train_dataloader = DataLoader(train_dataset, batch_size=params.batch_size, \
        shuffle=True, collate_fn=collate_fn_train)

    print(f'Step involved: {len(train_dataset)/params.batch_size}')
    t=time.time()

    for i, (clip, label, vid_path) in enumerate(train_dataloader):
        if i%10 == 0:
            print()
            clip = clip.permute(0,1,3,4,2)
            print(f'Full_clip shape is {clip.shape}')
            print(f'Label is {label}')
            # pickle.dump(clip, open('f_clip.pkl','wb'))
            # pickle.dump(label, open('label.pkl','wb'))
            # exit()
    print(f'Time taken to load data is {time.time()-t}')

    '''train_dataset = multi_baseline_dataloader_val_strong(split =1 , shuffle = False, data_percentage = 1.0,  mode = 2)
    train_dataloader = DataLoader(train_dataset, batch_size=params.batch_size, shuffle=True, collate_fn=collate_fn2)

    print(f'Step involved: {len(train_dataset)/params.batch_size}')
    t=time.time()

    for i, (clip, label, vid_path, frame_list) in enumerate(train_dataloader):
        if i%25 == 0:
            print()
            # clip = clip.permute(0,1,3,4,2)
            print(f'Full_clip shape is {clip.shape}')
            print(f'Label is {label}')
            # print(f'Frame list is {frame_list}')

            # pickle.dump(clip, open('f_clip.pkl','wb'))
            # pickle.dump(label, open('label.pkl','wb'))
            # exit()
    print(f'Time taken to load data is {time.time()-t}')'''