"""Tests for the config API."""

import unittest

from pygame_assets.exceptions import NoSuchConfigurationParameterError
from pygame_assets.configure import Config, ConfigMeta
from pygame_assets.configure import get_config, config_exists, remove_config
from pygame_assets.configure import get_environ_config, set_environ_config
from .utils import change_config, temp_config

class TestConfigMeta(unittest.TestCase):
    """Unit tests for the config._params API."""

    def test_get_class_attributes(self):
        class Foo:
            a = 1
            b = 2

            __builtins_are_ignored__ = 'bar'

            def functions_are_ignored():
                return 0

        self.assertDictEqual(ConfigMeta.get_class_attributes(Foo),
                             {'a': 1, 'b': 2})

    def test_two_config_instances_are_equal(self):
        config1 = Config()
        config2 = Config()
        self.assertEqual(config1, config2)

    def test_subconfig_inherits_config_meta(self):
        class SubConfig(Config):
            name = 'sub'

        config = Config()
        subclassed = SubConfig()
        for param in config._meta:
            self.assertIn(param, subclassed._meta)
            self.assertEqual(subclassed._meta[param], config._meta[param])

    def test_new_subconfig_parameters_are_saved(self):
        class SubConfig(Config):
            name = 'sub'

            class Meta:
                base = 'mybase'

        subclassed = SubConfig()
        self.assertEqual(subclassed._meta['base'], 'mybase')

    def test_subconfig_inherits_super_meta(self):

        with temp_config('custom', 'sub'):
            class SuperConfig(Config):
                name = 'custom'

                class Meta:
                    base = 'custom/base'

            class SubConfig(SuperConfig):
                name = 'sub'

            superconfig = SuperConfig()
            subconfig = SubConfig()
            self.assertEqual(subconfig.base, superconfig.base)

    def test_not_registered_parameter_raises_exception(self):
        with self.assertRaises(NoSuchConfigurationParameterError):
            class SubConfig(Config):
                name = 'sub'

                class Meta:
                    not_registered = 'will raise an exception'

    def test_config_has_meta_dict(self):
        config = get_config()
        self.assertIsNotNone(config._meta)
        self.assertIsInstance(config._meta, dict)

    def test_meta_contains_all_parameters(self):
        config = get_config()
        for param in ConfigMeta._allowed_params:
            self.assertIn(param, config._meta)

    def test_get_parameter_through_dot_notation(self):
        config = get_config()
        for param in config._meta:
            self.assertEqual(getattr(config, param), config._meta.get(param))

    def test_set_parameter_through_dot_notation(self):
        name = ConfigMeta._allowed_params[0]
        with change_config(name) as config:
            setattr(config, name, 'other')
            self.assertEqual('other', getattr(config, name))

    def test_set_parameter_propagates_accross_get_config(self):
        name = ConfigMeta._allowed_params[0]
        with change_config(name) as config:
            setattr(config, name, 'other')
            self.assertEqual(getattr(get_config(), name), 'other')

    def test_has_default_font_size(self):
        config = get_config()
        self.assertIsNotNone(config.default_font_size)

    def test_set_config_base(self):
        with change_config('base') as config:
            self.assertEqual('./assets', config.base)
            config.base = 'static'
            self.assertEqual('static', config.base)

class TestConfigureUtils(unittest.TestCase):
    """Test configuration management utilities."""

    def test_remove_config(self):
        class FooConfig(Config):
            name = 'foo'

        self.assertTrue(config_exists('foo'))
        remove_config('foo')
        self.assertFalse(config_exists('foo'))

    def test_config_exists(self):
        self.assertFalse(config_exists('foo'))

        class FooConfig(Config):
            name = 'foo'

        self.assertTrue(config_exists('foo'))
        remove_config('foo')

class TestConfigureFromEnv(unittest.TestCase):
    """Test the configuration from OS environment variables."""

    def setUp(self):
        set_environ_config(None)

    def test_set_environ_config(self):
        set_environ_config(None)
        self.assertIsNone(get_environ_config())
        set_environ_config('foo')
        self.assertEqual('foo', get_environ_config())

    def test_set_environ_config_to_none_is_explicit(self):
        with self.assertRaises(TypeError):
            set_environ_config()
        set_environ_config(None)

    def test_get_config_from_env_var(self):
        class MyConfig(Config):
            name = 'myconfig'

        set_environ_config('myconfig')
        config = get_config()
        self.assertEqual('myconfig', config.name)

class TestConfigDirs(unittest.TestCase):
    """Unit tests for the search directory configuration."""

    def test_config_dirs_is_dict_of_loader_names_and_dirs(self):
        config = get_config()
        for loader_name, dirs in config.dirs.items():
            self.assertIsInstance(loader_name, str)
            self.assertIsInstance(dirs, list)

    def test_search_dirs(self):
        with change_config('base', 'dirs') as config:
            config.base = 'assets'
            config.dirs['spritesheet'] = ['spritesheet', 'sheets']
            expected = ['assets/spritesheet', 'assets/sheets']
            actual = config.search_dirs('spritesheet')
            self.assertListEqual(expected, actual)

class TestConfigure(unittest.TestCase):
    """Unit tests for the Config API."""

    def setUp(self):
        set_environ_config(None)

    def test_load_default_config(self):
        explicit_config = get_config('default')
        self.assertEqual(explicit_config.name, 'default')
        no_args_config = get_config()
        self.assertEqual(no_args_config.name, 'default')

    def test_load_test_config(self):
        config = get_config('test')
        self.assertEqual(config.name, 'test')

    def test_new_config_must_have_name(self):
        with self.assertRaises(ValueError):
            class ConfigWithoutName(Config):
                name = ''

    def test_new_config_is_registered_to_get_config(self):
        # assert custom config does not exist already
        with self.assertRaises(KeyError):
            get_config('custom')

        class CustomConfig(Config):
            name = 'custom'
            base = 'custom/path/to/assets'

        # custom config should now be available through get_config()
        config = get_config('custom')
        self.assertEqual('custom', config.name)
        self.assertEqual('custom/path/to/assets', config.base)


from app.visual_detector.workers import (
    FrameReaderThread, TowerDetectorThread, ComponentDetectorThread, DefectDetectorThread,
    DefectTrackingThread, ResultsProcessorThread, TiltDetectorThread, DumperClassifierThread,
    WoodCracksDetectorThread
)
from app.visual_detector.defect_detectors import (
    TowerDetector, ComponentsDetector, DumperClassifier, WoodCrackSegmenter,
    DefectDetector, LineModifier, ConcretePoleHandler, WoodCracksDetector
)
from app.visual_detector.utils import (
    Drawer, ResultSaver, ObjectTracker
)
from typing import List
import queue
import os
import uuid
import time

class MainDetector:

    def __init__(
            self,
            save_path: str,
            batch_size: int,
            detectors_to_run: List[str],
            search_defects: bool = True,
            db=None
    ):
        # Path on the server where processed data gets stored
        self.save_path = save_path
        self.search_defects = search_defects
        # To keep track of file processing
        self.progress = dict()

        if search_defects:
            self.check_defects = True
            # Metadata extractor could be initialized here if required
        else:
            self.check_defects = False

        # Initialize Qs and workers
        try:
            # --- video / image processors ---
            self.files_to_process_Q = queue.Queue()
            self.frame_to_pole_detector = queue.Queue(maxsize=2)
            self.pole_to_comp_detector = queue.Queue(maxsize=2)
            self.comp_to_defect_detector = queue.Queue(maxsize=3)
            self.defect_det_to_defect_tracker = queue.Queue(maxsize=2)
            self.defect_tracker_to_writer = queue.Queue(maxsize=2)

            # --- dumper classifier ---
            self.to_dumper_classifier = queue.Queue(maxsize=3)
            self.from_dumper_classifier = queue.Queue(maxsize=3)
            # --- tilt detector ---
            self.to_tilt_detector = queue.Queue(maxsize=3)
            self.from_tilt_detector = queue.Queue(maxsize=3)
            # --- wood tower cracks ---
            self.to_wood_cracks_detector = queue.Queue(maxsize=3)
            self.from_wood_cracks_detector = queue.Queue(maxsize=3)
            # --- conrete tower cracks ---
            pass
        except Exception as e:
            print(f"Failed during queue initialization. Error: {e}")
            raise e

        # Initialize detectors and auxiliary modules
        try:
            # Object detectors
            self.pole_detector = TowerDetector()
            self.component_detector = ComponentsDetector()

            # Defect detectors
            self.defect_detector = DefectDetector(
                running_detectors=detectors_to_run,
                tilt_detector=(self.to_tilt_detector, self.from_tilt_detector),
                dumpers_defect_detector=(self.to_dumper_classifier, self.from_dumper_classifier),
                wood_crack_detector=(self.to_wood_cracks_detector, self.from_wood_cracks_detector),
                concrete_cracks_detector=None,
                insulators_defect_detector=None
            )

            # Auxiliary modules
            self.drawer = Drawer(save_path=save_path)
            self.result_saver = ResultSaver()
            self.object_tracker = ObjectTracker()

        except Exception as e:
            print(f"Failed during detectors initialization. Error: {e}")
            raise e

        # Initialize threads
        self.threads_to_run_and_join = list()
        try:
            self.frame_reader_thread = FrameReaderThread(
                batch_size=batch_size,
                in_queue=self.files_to_process_Q,
                out_queue=self.frame_to_pole_detector,
                progress=self.progress
            )
            self.threads_to_run_and_join.append(self.frame_reader_thread)

            self.pole_detector_thread = TowerDetectorThread(
                in_queue=self.frame_to_pole_detector,
                out_queue=self.pole_to_comp_detector,
                poles_detector=self.pole_detector,
                progress=self.progress
            )
            self.threads_to_run_and_join.append(self.pole_detector_thread)

            self.component_detector_thread = ComponentDetectorThread(
                in_queue=self.pole_to_comp_detector,
                out_queue=self.comp_to_defect_detector,
                component_detector=self.component_detector,
                progress=self.progress
            )
            self.threads_to_run_and_join.append(self.component_detector_thread)

            self.defect_detector_thread = DefectDetectorThread(
                in_queue=self.comp_to_defect_detector,
                out_queue=self.defect_det_to_defect_tracker,
                defect_detector=self.defect_detector,
                progress=self.progress,
                check_defects=search_defects
            )
            self.threads_to_run_and_join.append(self.defect_detector_thread)

            self.defect_tracker_thread = DefectTrackingThread(
                in_queue=self.defect_det_to_defect_tracker,
                out_queue=self.defect_tracker_to_writer,
                object_tracker=self.object_tracker,
                results_saver=self.result_saver,
                progress=self.progress,
                database=db
            )
            self.threads_to_run_and_join.append(self.defect_tracker_thread)

            self.results_processor_thread = ResultsProcessorThread(
                in_queue=self.defect_tracker_to_writer,
                save_path=save_path,
                drawer=self.drawer,
                progress=self.progress
            )
            self.threads_to_run_and_join.append(self.results_processor_thread)

            # -------------- DEFECT DETECTORS ---------------
            if "dumper" in detectors_to_run:
                dumper_classifier = DumperClassifier()
                self.dumper_classifier_thread = DumperClassifierThread(
                    in_queue=self.to_dumper_classifier,
                    out_queue=self.from_dumper_classifier,
                    dumper_classifier=dumper_classifier,
                    progress=self.progress
                )
                self.threads_to_run_and_join.append(self.dumper_classifier_thread)

            if "pillar" in detectors_to_run:
                self.tilt_detector_thread = TiltDetectorThread(
                    in_queue=self.to_tilt_detector,
                    out_queue=self.from_tilt_detector,
                    tilt_detector=ConcretePoleHandler(line_modifier=LineModifier),
                    progress=self.progress
                )
                self.threads_to_run_and_join.append(self.tilt_detector_thread)

            if "wood" in detectors_to_run:
                wood_tower_segment = WoodCrackSegmenter()
                self.wood_cracks_thread = WoodCracksDetectorThread(
                    in_queue=self.to_wood_cracks_detector,
                    out_queue=self.from_wood_cracks_detector,
                    wood_cracks_detector=WoodCracksDetector,
                    wood_tower_segment=wood_tower_segment,
                    progress=self.progress
                )
                self.threads_to_run_and_join.append(self.wood_cracks_thread)

        except Exception as e:
            print(f"Failed during workers initialization. Error: {e}")
            raise e

        # Start the app by launching all workers
        self.start()

    def predict(
            self,
            request_id: str,
            path_to_data: str,
            pole_number: int
    ) -> dict:
        """
        API endpoint - parses input data, puts files provided to the Q if of appropriate extension
        :param request_id:
        :param path_to_data:
        :param pole_number:
        :return:
        """
        file_IDS = {}

        if os.path.isfile(path_to_data):
            file_id = self.check_type_and_process(
                path_to_file=path_to_data,
                pole_number=pole_number,
                request_id=request_id
            )
            file_IDS[path_to_data] = file_id
            return file_IDS

        elif os.path.isdir(path_to_data):
            for item in os.listdir(path_to_data):
                path_to_file = os.path.join(path_to_data, item)
                file_id = self.check_type_and_process(
                    path_to_file=path_to_file,
                    pole_number=pole_number,
                    request_id=request_id
                )
                file_IDS[path_to_file] = file_id
        else:
            print(f"ERROR: Cannot process the file {path_to_data}, neither a folder nor a file")

        return file_IDS

    def check_type_and_process(
            self,
            path_to_file: str,
            pole_number: int,
            request_id: str
    ) -> str:
        """
        Checks whether a file can be processed
        :return: file's ID if it was put in the Q or "None" if extension is not supported
        """
        filename = os.path.basename(path_to_file)

        if any(filename.endswith(ext) for ext in ["jpg", "JPG", "jpeg", "JPEG", "png", "PNG"]):
            print(f"Added image {filename} to the processing queue")
            return self.process_file(
                path_to_file=path_to_file,
                pole_number=pole_number,
                file_type="image",
                request_id=request_id
            )

        elif any(filename.endswith(ext) for ext in ["avi", "AVI", "MP4", "mp4"]):
            print(f"Added video {filename} to the processing queue")
            return self.process_file(
                path_to_file=path_to_file,
                pole_number=pole_number,
                file_type="video",
                request_id=request_id
            )
        else:
            print(f"ERROR: file {filename}'s extension is not supported")
            return "None"

    def process_file(
            self,
            path_to_file: str,
            pole_number: int,
            file_type: str,
            request_id: str
    ) -> str:
        """

        :param path_to_file:
        :param pole_number:
        :param file_type:
        :param request_id:
        :return:
        """
        # Each file to process gets a unique ID number to track its progress
        file_id = str(uuid.uuid4())

        # Keep track of processing progress
        self.progress[file_id] = {
            "status": "Awaiting processing",
            "request_id": request_id,
            "file_type": file_type,
            "pole_number": pole_number,
            "path_to_file": path_to_file,
            "total_frames": None,
            "now_processing": 0,
            "processed": 0,
            "defects": []
        }
        self.files_to_process_Q.put(file_id)

        return file_id

    def start(self) -> None:
        for thread in self.threads_to_run_and_join:
            thread.start()

    def stop(self) -> None:
        """Wait until all threads complete their job and exit"""
        self.files_to_process_Q.put("STOP")
        for thread in self.threads_to_run_and_join:
            thread.join()


import pygame

class Heart(pygame.sprite.Sprite):
    def __init__(self, imgfile, hp, pos=[]):
        super().__init__()

        self.heart_img = pygame.image.load(imgfile)
        self.rect = self.heart_img.get_rect()
        self.pos = pos
        self.HP = hp
        self.speed = 7
        self.Invincible = 0 #受到攻击后的无敌时间

    # 裂开的特效（待完善） ；return 是否活着
    def is_dead(self):
        if self.HP <= 0:
            return True
        else:
            return False

    # 居中画图
    def target_pos(self, pos):
            pos = (pos[0]-self.rect[-2]//2,pos[1]-self.rect[-1]//2)
            return pos


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
Conjunto de funciones desarrolladas para ser utilizadas en los programas principales
'''
import numpy as np
import torch
import matplotlib.pyplot as plt

from labcomdig import gray2de
from network import *

'''
NUMPY
'''
def transmisorpamV2(Bn,Eb,M,p,L):
    """
    [Xn,Bn,An,phi,alfabetopam] = transmisorpamV2(Bn,Eb,M,p,L)

    Entradas:
     Bn = Secuencia de dígitos binarios
     Eb = Energía media por bit transmitida en Julios
     M  = Número de síímbolos del código PAM
     p  = Pulso paso de baja o paso de banda

     L  = Número de puntos a utilizar en la representación de un símbolo

    Devuelve:
     Xn = la señal de información (discreta)
     Bn = La secuencia de dígitos binarios realmente transmitidos
     An = La secuencia de niveles de amplitud transmitidos
   	 phi = Pulso básico real normalizado (energía unidad)
     alfabetopam = Los niveles de amplitud asociados a cada símbolo
    """
    # Comprobación de parámetros de entrada
    p=np.reshape(p,[np.size(p)])  # <=   #p=p.squeeze()
    if len(Bn)<1 or Eb<=0 or M<2 or np.dot(p,p)==0 or L<1:
        raise Exception('Error: revise los parámetros de entrada')
    # Se obtienen en primer lugar los niveles asociado a cada símbolo ¿Cuántos bits hay en cada símbolo?
    k = int(np.ceil(np.log2(M)))
    M = 2**(k) # Se Ajusta M a una potencia de dos

    # El alfabeto [Ver la ecuación (4.21)]
    alfabetopam = np.sqrt(3*Eb*np.log2(M)/(M**2-1))*(2*(np.arange(M))-M+1)

    # Si Bn no tiene una longitud múltiplo de k, se completa con ceros
    Nb = len(Bn)  # Número de bits a transmitir, actualizado
    Bn = Bn.squeeze().astype(int) #Comprobación de int y dimensiones
    Bn = np.r_[Bn,np.zeros(int(k*np.ceil(Nb/k)-Nb)).astype(int)] #
    Nb = len(Bn)  # Número de bits a transmitir tras la corrección
    Ns = Nb//k # Número de símbolos a transmitir

    # La secuencia generada
    if M>2:
        An = alfabetopam[gray2de(np.reshape(Bn,[Ns,k]))]
    else:
        An = alfabetopam[Bn]

    # Comprobación de las longitudes y otros datos del pulso suministrado para
    # hacer que el número de muestras del mismo sea efectivamente L
    Ls = len(p)
    if Ls<L:
        p = np.r_[p, np.zeros(L-Ls)]
    elif Ls>L:
        print('La duración del pulso se trunca a {} muestras'.format(str(L)))
        p = p[:L] #Debe modificarse si se quiere un pulso de más de L muestras
    # Se normaliza la energía del pulso para obtener la base del sistema
    phi = p / np.sqrt(p@p)

    # Obtención del tren de pulsos, Xn = np.kron(An,phi) ó
    Xn = np.reshape(np.reshape(An,[Ns,1])*phi,[Ns*L,]) #Debe modificarse si se quiere un pulso de más de L muestras
    return [Xn,Bn,An,phi,alfabetopam]

def de2Nary(d,n,N):
    """
    b = de2Nary(d,n,N)
    Convierte un número decimal, d, en un vector binario, b, de longitud n
    con base N
    """
    c = np.zeros([len(d),int(n)])
    for i in range(int(n)): d, c[:,i] = np.divmod(d,N)
    c = np.fliplr(c);
    return c.astype(int).T

'''
TORCH
'''

def t_Qfunct(x):
    """
     y = Qfunct(x) evalúa la función Q en x.
    Donde y = 1/sqrt(2*pi) * integral desde x hasta inf de exp(-t^2/2) dt
    """
    y=(1/2)*torch.special.erfc(x/(torch.tensor(2)**.5))
    return y

def t_gray2de(b):
    """
    Convierte cada columna de la matriz formada por dígitos binarios b en un vector
        fila de los valores decimales correspondientes, aplicando codificación de Gray.
    """
    if not torch.is_tensor(b):
        raise Exception('Error: la entrada no es un tensor')
    b = b.long() #Aseguro que sea tipo int
    c = torch.zeros_like(b); c[:,0] = b[:,0]
    for i in range(1,b.size(dim=1)):
        c[:,i] = torch.logical_xor(c[:,i-1], b[:,i])
    c = torch.fliplr(c) # Convierte los bits menos significativos en los más significativos
    #Comprueba un caso especial.
    [n,m] = c.size()
    if torch.min(torch.tensor((m,n))) < 1:
        d = []
        return
    d = c @ 2**torch.arange(m)
    return d

def t_de2gray(d,n):
    """
    b = de2gray(d,n)
    Convierte un número decimal, d, en un vector binario, b, de longitud n
    """
    c = torch.zeros(len(d),int(n))
    for i in range(int(n)):
        c[:,i] = torch.fmod(d,2) #resto
        d = torch.div(d,2,rounding_mode='floor') #cociente
    c = torch.fliplr(c); b = torch.zeros_like(c); b[:,0] = c[:,0]; aux = b[:,0]
    for i in range(1,int(n)):
        b[:,i] = torch.logical_xor(aux, c[:,i])
        aux = torch.logical_xor(b[:,i], aux)
    return torch.reshape(b,[-1]).long()

def t_simbolobit(An,alfabeto):
    """
    Bn = simbolobit(An, alfabeto)
    An = secuencia de símbolos pertenecientes al alfabeto
    alfabeto = tabla con los símbolos utilizados en la transmisión
    Bn = una secuencia de bit, considerando que los símbolos se habían
    generado siguiendo una codificación de Gray
    """

    k = torch.log2(torch.tensor(len(alfabeto))) # bits por símbolo

    if k>1:
        distancia = abs(alfabeto[0]-alfabeto[1])
        indices   = torch.round((An-alfabeto[0])/distancia)
        #Bn        = torch.reshape(t_de2gray(indices,k),int(k*len(An)))
        Bn = t_de2gray(indices,k)
    else:
        Bn = ((An/max(alfabeto))+1)/2

    return Bn

def t_detectaSBF(r,alfabeto):
    """
    An = detectaSBF(r,alfabeto)
    r = secuencia a la entrada del detector, con estimaciones de s_i
    alfabeto = tabla con los niveles de amplitud/símbolos

    Genera:
    An = una secuencia de símbolos pertenecientes al alfabeto de acuerdo con
    una regla de distancia euclidiana mínima (mínima distancia)
    """

    # Obtiene el índice respecto al alfabeto
    r_repeated = torch.repeat_interleave(r.reshape(r.size(dim=0),1),alfabeto.size(dim=0),1)
    ind = torch.argmin(torch.abs(r_repeated - alfabeto), 1)

    # Genera la secuencia de niveles detectados
    An = alfabeto[ind]

    return An

def t_transmisorpamV2(Bn,Eb,M,p,L):
    """
    [Xn,Bn,An,phi,alfabetopam] = transmisorpamV2(Bn,Eb,M,p,L)

    Entradas:
     Bn = Secuencia de dígitos binarios
     Eb = Energía media por bit transmitida en Julios
     M  = Número de síímbolos del código PAM
     p  = Pulso paso de baja o paso de banda
     L  = Número de puntos a utilizar en la representación de un símbolo

    Devuelve:
     Xn = la señal de información (discreta)
     Bn = La secuencia de dígitos binarios realmente transmitidos
     An = La secuencia de niveles de amplitud transmitidos
   	 phi = Pulso básico real normalizado (energía unidad)
     alfabetopam = Los niveles de amplitud asociados a cada símbolo
    """
    # Paso a tensores de los parámetros de entrada
    Eb = torch.tensor(Eb)
    M = torch.tensor(M)
    L = torch.tensor(L)

    # Comprobación de parámetros de entrada
    p=torch.reshape(p,p.size())  # <=   #p=p.squeeze()
    if Bn.size(dim=1)<1 or Eb<=0 or M<2 or torch.dot(p,p)==0 or L<1:
        raise Exception('Error: revise los parámetros de entrada')
    # Se obtienen en primer lugar los niveles asociado a cada símbolo ¿Cuántos bits hay en cada símbolo?
    k = torch.ceil(torch.log2(M))
    M = 2**(k) # Se Ajusta M a una potencia de dos

    # El alfabeto [Ver la ecuación (4.21)]
    alfabetopam = torch.sqrt(3*Eb*torch.log2(M)/(M**2-1))*(2*(torch.arange(M))-M+1)

    # Si Bn no tiene una longitud múltiplo de k, se completa con ceros
    Nb = torch.tensor(Bn.size(dim=1))  # Número de bits a transmitir, actualizado
    Bn = Bn.squeeze().long() #Comprobación de int y dimensiones
    Bn = torch.hstack((Bn,torch.zeros((k*torch.ceil(Nb/k)-Nb).long()))) #
    Nb = torch.tensor(Bn.size(dim=0))  # Número de bits a transmitir tras la corrección
    Ns = torch.div(Nb,k,rounding_mode='trunc') # Número de símbolos a transmitir

    # La secuencia generada
    if M>2:
        An = alfabetopam[t_gray2de(torch.reshape(Bn,(int(Ns),int(k))))]
    else:
        An = alfabetopam[Bn.long()]

    # Comprobación de las longitudes y otros datos del pulso suministrado para
    # hacer que el número de muestras del mismo sea efectivamente L
    Ls = p.size(dim=0)
    if Ls<L:
        p = torch.hstack(p,torch.zeros(L-Ls))
    elif Ls>L:
        print('La duración del pulso se trunca a {} muestras'.format(str(L)))
        p = p[:L] #Debe modificarse si se quiere un pulso de más de L muestras
    # Se normaliza la energía del pulso para obtener la base del sistema
    phi = p / torch.sqrt(p@p)

    # Obtención del tren de pulsos, Xn = np.kron(An,phi) ó
    a = torch.reshape(An,(int(Ns),1))*phi
    Xn = torch.reshape(a,(int(Ns)*L,)) #Debe modificarse si se quiere un pulso de más de L muestras
    return [Xn,Bn.long(),An,phi,alfabetopam]

def t_de2Nary(d,n,N):
    """
    b = de2Nary(d,n,N)
    Convierte un número decimal, d, en un vector binario, b, de longitud n
    con base N
    """
    c = torch.zeros(len(d),int(n))
    for i in range(int(n)):
        c[:,i] = torch.fmod(d,N) #resto
        d = torch.div(d,N,rounding_mode='floor') #cociente
    c = torch.fliplr(c);
    return c.long().T

'''
NEURAL NETWORK
'''

def correction(Nb_train,Nb_test,k,M,n):
    """
    [Nb_train,Ns_train,Nb_test,Ns_test] = correction(Nb_train,Nb_test,k,M,n)

    Cambia el valor de Nb_train y Nb_test si es necesario para
    que sea múltiplo de k y M**n
    """

    Ns_train = int(torch.floor(torch.tensor(Nb_train/(k*M**n)))) #Number of symbols (multiple of k and M**n)
    Nb_train = int(Ns_train*(k*M**n)) #Number of bits on transmission
    Ns_train = int(torch.floor(torch.tensor(Nb_train/k))) #Number of symbols (multiple of k)

    Ns_test = int(torch.floor(torch.tensor(Nb_test/k))) #Number of symbols (multiple of k)
    Nb_test = int(Ns_test*k)          #Number of bits on transmission

    return [Nb_train,Ns_train,Nb_test,Ns_test]

def create_datasets(H,Nb_train,Nb_test,Eb,M,n,SNRdb,batch_size,valid_size):
    """
    [trainloader,validloader,testloader,x,alphabet] = create_datasets(H,Nb_train,Nb_test,Eb,M,n,SNRdb,batch_size,valid_size)

    Devuelve los conjuntos de datos listos para itererar, además del conjunto
    de combinaciones posibles y el alfabeto de la constelación
    """
    # Data generation for training
    [Rn,Cn,Bn,x,alphabet] = generate_data(H,Nb_train,Eb,M,n,SNRdb,train=True)
    trainset = SymbolsDataset(Rn,Cn,Bn)

    # Data generation for testing
    [Rn,Cn,Bn,x,alphabet] = generate_data(H,Nb_test,Eb,M,n,SNRdb,train=False)
    testset = SymbolsDataset(Rn,Cn,Bn)

    # Indixes used for validation
    num_train = len(trainset)
    split = int(np.floor(valid_size * num_train))

    # Split data
    train_split, valid_split = SymbolsDataset.split_data(trainset,split)

    # Load training data in batches
    trainloader = torch.utils.data.DataLoader(train_split,
                                              batch_size=batch_size,
                                              num_workers=0)

    # Load validation data in batches
    validloader = torch.utils.data.DataLoader(valid_split,
                                              batch_size=batch_size,
                                              num_workers=0)

    # Load test data in batches
    testloader = torch.utils.data.DataLoader(testset,
                                             batch_size=batch_size,
                                             num_workers=0)

    return [trainloader,validloader,testloader,x,alphabet]

def generate_data(H,Nb,Eb,M,n,SNRdb,train):
    """
    [Xn/Rn,Bn,Cn,x,alphabet] = generate_data(H,Nb,Eb,M,n,SNRdb,test)

    Entradas:
     H     = Matriz del canal
     Nb    = Número de bits a transmitir
     Eb    = Energía media por bit transmitida en Julios
     M     = Número de símbolos del código PAM
     n     = Número de antenas del sistema MIMO nxn
     SNRdb = Signal to Noise Ratio
     test  = Booleano para determinar si son datos para entrenar o testar

    Devuelve:
     Rn = La señal de información (discreta) recibida
     Bn = La secuencia de dígitos binarios realmente transmitidos
     Cn = La combinación correspondiente a cada valor de An
     x  = La matriz con las posibles combinaciones
     alphabet = Los niveles de amplitud asociados a cada símbolo
    """

    k = int(torch.log2(torch.tensor(M)))      #Number of bits on each symbol
    Ns = int(torch.floor(torch.tensor(Nb/k))) #Number of symbols (multiple of k)

    alphabet = torch.sqrt(3*Eb*torch.log2(torch.tensor(M))/(M**2-1))*(2*(torch.arange(M))-M+1)

    # possible combinations
    ind = t_de2Nary(torch.arange(M**n),n,M)
    x = alphabet[ind]

    if train:  # data for training. same number of symbols transmitted

        Xn = x.tile(int(Ns/x.size(1)))
        Xn = Xn[:,torch.randperm(Xn.size(1))] # Shuffle
        Bn = t_simbolobit(Xn.flatten(),alphabet) #Detected bits

        #Data reshape. Each row represents an antenna
        Bn = Bn.reshape(n,Nb)

        #Index of the combinations
        Cn = getidx(x,Xn,n)

    else:     # data for testing

        # fixed seed for same results always
        torch.manual_seed(1)

        bn = torch.randint(0,2,(1,Nb*n)) #Bits to transmit
        [Xn,Bn,An,phi,alphabet] = t_transmisorpamV2(bn,Eb,M,torch.ones(1),1)

        #Data reshape. Each row represents an antenna
        Xn = Xn.reshape(n,Ns)
        Bn = Bn.reshape(n,Nb)
        An = An.reshape(n,Ns)

        #Index of the combinations
        Cn = getidx(x,An,n)

    SNR = 10**(SNRdb/10)   #Signal to Noise Ratio [n.u.]
    varzn = Eb/(2*SNR)         #Noise variance
    #if varzn <= 0.0025: varzn = torch.tensor(0.0025)
    Wn = torch.sqrt(varzn)*torch.randn(*Xn.shape)  #AWGN

    Rn = H@Xn + Wn

    return [Rn,Cn,Bn.long(),x,alphabet]

def getidx(x,An,n):
    """
    idx = getidx(x,An,n)

    Función que retorna el número de combinación correspondiente a cada
    valor de An

    x  = La matriz con las posibles combinaciones
    An = La secuencia de niveles de amplitud transmitidos
    n  = El número de antenas
    """

    idx = torch.empty(0)
    for col in An.T:
        i = (sum(x==col.view(n,1)) == n).nonzero(as_tuple=True)[0]
        idx = torch.hstack((idx,i))
    return idx.long()

def train_model(model, trainloader, validloader, optimizer, criterion, patience, n_epochs):
    """
    [model, avg_train_losses, avg_valid_losses] = create_datasets(H,Nb_train,Nb_test,Eb,M,n,SNRdb,batch_size,valid_size):

    Función que entrena a la red y retorna el modelo y los errores cometidos

    """
    min_valid_loss = np.inf
    es_counter = 0

    # loss per batch
    train_losses = []
    valid_losses = []

    # loss per epoch
    avg_train_losses = []
    avg_valid_losses = []

    for epoch in range(1, n_epochs + 1):

        # Train the model
        model.train()
        for symbols, combs, bits in trainloader:
            # clear the gradients
            optimizer.zero_grad()
            # forward pass
            output = model(symbols)
            # calculate the loss
            loss = criterion(output, combs)
            # backward  pass
            loss.backward()
            # perform a single optimization step
            optimizer.step()
            # record training loss
            train_losses.append(loss.item())

        # Validate the model
        model.eval()
        for symbols, combs, bits in validloader:
            # forward pass
            output = model(symbols)
            # calculate the loss
            loss = criterion(output, combs)
            # record valid loss
            valid_losses.append(loss.item())

        # calculate average loss over an epoch
        train_loss = np.average(train_losses)
        valid_loss = np.average(valid_losses)
        avg_train_losses.append(train_loss)
        avg_valid_losses.append(valid_loss)

        if min_valid_loss > valid_loss:
            min_valid_loss = valid_loss
            # Saving State Dict
            torch.save(model.state_dict(), 'saved_model.pth')
        else:
            es_counter += 1
            if es_counter == patience: break

        # load saved model
        model.load_state_dict(torch.load('saved_model.pth'))

        # clear lists to track next epoch
        train_losses = []
        valid_losses = []

    return [model, avg_train_losses, avg_valid_losses]

def printloss(avg_train_losses,avg_valid_losses,SNR,save):
    """
    Función para mostrar por consola las pérdidas de entramiento y validación
    en un epoch
    """

    fig = plt.figure(1,figsize=(10,8))
    plt.plot(range(1,len(avg_train_losses)+1),avg_train_losses, label='Training Loss')
    plt.plot(range(1,len(avg_valid_losses)+1),avg_valid_losses,label='Validation Loss')

    # find position of lowest validation loss
    minposs = avg_valid_losses.index(min(avg_valid_losses))+1
    plt.axvline(minposs, linestyle='--', color='r',label='Early Stopping Checkpoint')

    plt.xlabel('epochs')
    plt.ylabel('loss')
    plt.ylim(0, 2) # VARIABLE scale
    plt.xlim(0, len(avg_valid_losses)+1) # consistent scale
    plt.grid(True)
    plt.legend()
    plt.title('SNR = {}dB'.format(SNR))
    plt.tight_layout()
    plt.show()
    if save: fig.savefig('loss_plot.png', bbox_inches='tight')
    pass

def eval_model(model, trainloader, validloader, testloader, x, alphabet, k):
    """
    [berTrain, berValid, berTest] = eval_model(model, trainloader, validloader, testloader, x, alphabet, k)

    Función que evalua el modelo y devuelve las distintas tasas de error binarias
    """

    # BER for saved model with train data
    model.eval()

    berTrain = 0
    for symbols, combs, bits in trainloader:
        # forward pass
        output = model(symbols)
        berTrain += ber(output,x,alphabet,bits,k)
    berTrain /= len(trainloader)

    berValid = 0
    for symbols, combs, bits in validloader:
        # forward pass
        output = model(symbols)
        berValid += ber(output,x,alphabet,bits,k)
    berValid /= len(validloader)

    berTest = 0
    for symbols, combs, bits in testloader:
        # forward pass
        output = model(symbols)
        berTest += ber(output,x,alphabet,bits,k)
    berTest /= len(testloader)

    return [berTrain, berValid, berTest]

def ber(output,x,alphabet,bits,k):
    """
    ber = ber(output,x,alphabet,bits,k)
    Calcula la tasa de error de bit de una señal transmitida

    Entradas:
     output   = Salida de la Red Neuronal (bacth_size x n_combinations)
     x        = La matriz con las posibles combinaciones
     alphabet = Los niveles de amplitud asociados a cada símbolo
     bits     = El conjunto de Bn transmitidos
     k        = Los bits que hay por símbolo
    """

    Nb = k*output.size(0) # Número de bits transmitidos (k*Ns)
    n = x.size(0)         # Núermo de antenas

    val, idx = torch.max(output,1) # El máximo de cada salida

    # Es necesario modificar la variable 'bits' de manera que estén
    # en la forma nxNb, correspondiendo cada fila a la transmisión de una antena
    Bn = torch.empty(0,Nb)
    for i in range(int(n)): Bn = torch.vstack((Bn,bits[:,i,:].flatten()))

    Andetected = x[:,idx]    # Detected symbols
    Bndetected = t_simbolobit(Andetected.flatten(),alphabet).reshape(n,Nb) #Detected bits
    errorsBit = Nb-torch.sum(Bndetected==Bn,axis=1)

    return (errorsBit/Nb).reshape(n,1)

def detect(Rn,Bn,H,alphabet,Nb,Ns,n,method,**kwargs):
    """
    ber = detect(Rn,Bn,H,alphabet,Nb,Ns,n,method,**kwargs)

    Calcula la BER de los métodos ZF, LMMSE y ML
    """
    if method == 'ZF':
        Hinv = kwargs.get('Hinv',None)

        Z = Hinv@Rn

    elif method == 'LMMSE':
        Es = kwargs.get('Es',None)
        varzn = kwargs.get('varzn',None)

        Z = torch.linalg.inv(H.T@H + varzn/Es*torch.eye(n))@H.T@Rn

    elif method == 'ML':
        x = kwargs.get('x',None)

        hx = torch.unsqueeze(H@x,2) #H@x with shape (n,combinations,1)
        Z = x[:,torch.argmin(torch.linalg.norm(Rn.reshape(n,1,Ns)-hx,dim=0),axis=0)]

    Andetected = t_detectaSBF(Z.flatten(),alphabet).reshape(n,Ns) #Detected symbols
    Bndetected = t_simbolobit(Andetected.flatten(),alphabet).reshape(n,Nb) #Detected bits
    errorsBit = Nb-torch.sum(Bndetected==Bn,axis=1) # Nº Errors per bit

    return (errorsBit/Nb).reshape(n,1)


import argparse
from dis import dis
import gym
import numpy as np
from itertools import count

import pyximport; pyximport.install()

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import MultivariateNormal
import matplotlib.pyplot as plt
import pybullet
import os

from tensorboardX import SummaryWriter

from config import env_name, log_interval, device, lr, gamma, env_seed, torch_seed, rand_init

from arguments import get_args

def save_state(net, opt, torch_seed, env_seed, fname, np_seed=None, py_seed=None,):
    # save both model state and optimizer state
    states = {
        'state_dict': net.state_dict(),
        'optimizer': opt.state_dict(),
        'torch_seed': torch_seed,
        'env_seed': env_seed,
        'gamma': gamma,
        'lr' : lr
        # 'py_seed': py_seed,
        # 'np_seed': np_seed
    }
    torch.save(states, fname)

class GaussNet(nn.Module):
    def __init__(self):
        super(GaussNet, self).__init__()
        self.affine1 = nn.Linear(8, 32)
        self.linear = nn.Linear(32,32)
        self.affine2 = nn.Linear(32, 2)

    def forward(self, x):
        x = self.affine1(x)
        x = F.relu(x)
        x = self.linear(x)
        x = F.relu(x)
        mu = self.affine2(x)
        x = F.relu(x)
        return mu

def get_action(state, net, device):
    state = torch.from_numpy(state).float().to(device)
    mu = net(state).to(device)
    #calculate the probability
    dist = MultivariateNormal(mu, 0.1*torch.eye(2).to(device))
    action = dist.sample()
    log_prob = dist.log_prob(action)
    return action, log_prob

def update_policy(policy_loss,optimizer, device):
    optimizer.zero_grad()
    policy_loss = torch.cat(policy_loss).sum().to(device)
    policy_loss.backward()
    optimizer.step()

def main():
    args = get_args()
    gym.logger.set_level(40)
    env = gym.make(env_name,rand_init=rand_init)
    env.seed(env_seed)
    torch.manual_seed(torch_seed)

    state= env.reset()

    num_inputs = state.shape[0]
    num_actions = env.action_space.shape[0]
    print('state size:', num_inputs)
    print('action space shape, highs, lows : ', env.action_space.shape," ",env.action_space.high," ",env.action_space.low)

    net = GaussNet().to(device)
    optimizer = optim.Adam(net.parameters(), lr=lr)

    writer = SummaryWriter('logs')

    if type(args.load_train) == list and (args.load_train[0] >-1 and args.load_train[1] >-1) :
        load_path='{}_ep_{}_it_{}_type_{}.pkl'.format(args.rl_alg,args.load_train[0],args.load_train[1],args.alg_type)
        if not os.path.exists(os.path.join(args.model_path,load_path)):
            raise Exception("load path {} doesn't exist".format(load_path))
        os.system('clear')
        print(" Loading ... {}".format(load_path))

        net.load_state_dict(torch.load(os.path.join(args.model_path,load_path))["state_dict"])
        optimizer.load_state_dict(torch.load(os.path.join(args.model_path,load_path))["optimizer"])
        start_iteration = args.load_train[1]
        start_episode = 0
    else:
        start_iteration = 0
        start_episode = 0
    net.train()

    steps=0

    avg_itr_reward = []
    num_iter = args.num_iterations
    num_episodes = args.num_episodes
    for itr in range(start_iteration,start_iteration + num_iter):
        total_ep_reward = 0
        log_probs_itr = []
        policy_loss_itr = []
        returns_itr = []
        for i_episode in range(num_episodes):
            state, ep_reward = env.reset(), 0
            done = False
            rewards = []
            log_probs = []
            returns=[]
            while not done:
                steps=steps+1;
                action, log_prob = get_action(state, net, device)
                state, reward, done, _ = env.step(action.cpu().numpy())
                # print(reward)
                ep_reward +=reward
                rewards.append(reward)
                log_probs.append(log_prob)
            #Compute the returns from the rewards
            R = 0
            for r in rewards[::-1]:
                R = r + gamma * R
                returns.insert(0, R)
            returns = torch.FloatTensor(returns).to(device)
            #Store the log probability and returns to calculate the loss function
            log_probs_itr.append(log_probs)
            returns_itr.append(returns)
            total_ep_reward += ep_reward

        returns_ep = torch.cat(returns_itr)
        # Find the mean and std of the returns over the episodes to compute
        mean  = returns_ep.mean()
        std = returns_ep.std()
        for log_probs, returns in zip(log_probs_itr,returns_itr):
            # Use VANILLA REINFORCE with baseline
            returns = (returns - mean) / std
            policy_loss_ep = []
            for log_prob, return_ in zip(log_probs,returns):
                policy_loss_ep.append(-log_prob * return_)
            policy_loss_ep = torch.stack(policy_loss_ep).to(device)
            policy_loss_itr.append(policy_loss_ep)
        avg_itr_reward.append(total_ep_reward / num_episodes)
        # Update the policy neural network parameter using the Loss function
        update_policy(policy_loss_itr,optimizer,device)
        if itr % log_interval == 0:
            print('iteration: {} | loss: {:.3f} | steps: {} | avg_reward: {:.2f}'.format(
                itr, policy_loss_ep.sum(), steps, avg_itr_reward[-1]))
            writer.add_scalar('log/avg_reward', float(avg_itr_reward[-1]), itr)
            writer.add_scalar('log/loss', float(policy_loss_ep.sum()), itr)
    if type(args.load_train) == list and (args.load_train[0] >-1 and args.load_train[1] >-1) :
        model_path='{}_ep_{}_it_{}_type_{}.pkl'.format(args.rl_alg,args.num_episodes,args.load_train[1]+args.num_iterations,args.alg_type)
    else:
        model_path='{}_ep_{}_it_{}_type_{}.pkl'.format(args.rl_alg,args.num_episodes,args.num_iterations,args.alg_type)

    if not os.path.exists(args.model_path):
        os.makedirs(args.model_path)
    save_state(net, optimizer, torch_seed, env_seed, os.path.join(args.model_path,model_path))

    plt.figure()
    plt.plot(range(itr+1),avg_itr_reward)
    plt.title("Average Cumulative Reward vs iteration")
    plt.xlabel("iteration")
    plt.ylabel("Average Reward")
    if not os.path.exists(args.save_path):
        os.makedirs(args.save_path)
    plt.savefig("{}/{}_plot.png".format(args.save_path,model_path[:-4]))

if __name__ == '__main__':
    main()


'''
This code is based on https://github.com/ekwebb/fNRI which in turn is based on https://github.com/ethanfetaya/NRI
(MIT licence)
'''
from synthetic_sim_comperrors import *
import time
import numpy as np
import argparse
import os
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt

parser = argparse.ArgumentParser()
parser.add_argument('--num-train', type=int, default=100,
                    help='Number of training simulations to generate.')
parser.add_argument('--num-valid', type=int, default=10000,
                    help='Number of validation simulations to generate.')
parser.add_argument('--num-test', type=int, default=10000,
                    help='Number of test simulations to generate.')
parser.add_argument('--length', type=int, default=10000,
                    help='Length of trajectory.')
parser.add_argument('--length-test', type=int, default=10000,
                    help='Length of test set trajectory.')
parser.add_argument('--sample-freq', type=int, default=100,
                    help='How often to sample the trajectory.')
parser.add_argument('--n-balls', type=int, default=5,
                    help='Number of balls in the simulation.')
parser.add_argument('--seed', type=int, default=42,
                    help='Random seed.')
parser.add_argument('--savefolder', type=str, default='springcharge_comperrors',
                    help='name of folder to save everything in')
parser.add_argument('--sim-type', type=str, default='springcharge',
                    help='Type of simulation system')
parser.add_argument('--timesteptest', type=str, default=True,
                    help='Generate many datasets with different timesteps')
parser.add_argument('--animation', type=str, default=False,
                    help='Generate animation of many datasets with different timesteps')

args = parser.parse_args()
os.makedirs(args.savefolder)
par_file = open(os.path.join(args.savefolder,'sim_args.txt'),'w')
print(args, file=par_file)
par_file.flush()
par_file.close()

# no noise-useless for our investigation
if args.sim_type == 'springcharge':
    sim = SpringChargeSim(noise_var=0.0, n_balls=args.n_balls, box_size=5.0)

elif args.sim_type == 'springchargequad':
    sim = SpringChargeQuadSim(noise_var=0.0, n_balls=args.n_balls, box_size=5.0)

elif args.sim_type == 'springquad':
    sim = SpringQuadSim(noise_var=0.0, n_balls=args.n_balls, box_size=5.0)

elif args.sim_type == 'springchargefspring':
    sim = SpringChargeFspringSim(noise_var=0.0, n_balls=args.n_balls, box_size=5.0)

np.random.seed(args.seed)

def generate_mse_pertime(num_sims, length, sample_freq):
    loc_all = []
    vel_all = []
    _delta_T_ = []
    for i in range(num_sims):
        t = time.time()
        loc_mse, vel_mse, delta_T = sim.sample_trajectory(T=length, sample_freq=sample_freq)
        # dim 0 is the batch dimension, dim 1 is time: already averaged MSE over particles
        loc_all.append(loc_mse)
        vel_all.append(vel_mse)
        _delta_T_ = delta_T
    mse_loc = (np.asarray(loc_all).mean(axis=0)) / num_sims
    mse_vel = (np.asarray(vel_all).mean(axis=0)) / num_sims

    return mse_loc, mse_vel, delta_T

if (args.timesteptest):
    print("Calculating MSE over time due to computational errors")
    mse_loc, mse_vel, delta_T = generate_mse_pertime(args.num_train, args.length, args.sample_freq)
    # dim 0 on mse is different delta_T, dim 1 is different times along the motion, then (x,y)
    # [0,:,:] is the delta_T that the model gets
    mse_model = mse_loc[0].mean(axis=1)
    mse_model_vel = mse_vel[0].mean(axis=1)
    np.save(os.path.join(args.savefolder, 'mse_model_pos.npy'), mse_model)
    np.save(os.path.join(args.savefolder, 'mse_model_vel.npy'), mse_model_vel)
    mse_loc_mean = (np.asarray(mse_loc).mean(axis=1).mean(axis=1))
    fig = plt.figure()
    plt.plot(delta_T, mse_loc_mean)
    plt.xlabel('Time-step of Integration/(arbitrary units)')
    plt.ylabel('Averaged Mean Square Error/(arbitrary units)')
    plt.show()
    """
    https://jakevdp.github.io/PythonDataScienceHandbook/04.12-three-dimensional-plotting.html
    """
    y = np.multiply(delta_T, 10.0)
    x = np.multiply(np.arange(0,(int(args.length / args.sample_freq)-1)*0.001, 0.001), 10.0)
    Z = np.multiply((np.asarray(mse_loc).mean(axis=2)), 100.0)
    X, Y = np.meshgrid(x,y)
    fig = plt.figure()
    ax = plt.axes(projection = '3d')
    ax.plot_surface(X,Y,Z, cmap = 'plasma', edgecolor= 'none')
    ax.set_ylabel('Time-step of Integration')
    ax.set_xlabel('Time along trajectory')
    ax.set_zlabel('Averaged Mean Square Error')
    ax.view_init(45, -35)
    fig.savefig(os.path.join(args.savefolder, 'timestepmse.eps'), format='eps')
    # fig.savefig(os.path.join(args.savefolder, 'timestepmse.png'))
    ax.view_init(75,-35)
    fig.savefig(os.path.join(args.savefolder, 'timestepmse_birdseye.eps'), format='eps')
    # fig.savefig(os.path.join(args.savefolder, 'timestepmse_birdseye.png'))
    ax.view_init(45,225)
    fig.savefig(os.path.join(args.savefolder, 'timestepmse_rotated.eps'), format='eps')
    # fig.savefig(os.path.join(args.savefolder, 'timestepmse_rotated.png'))
    plt.show()
    # for i in range(len(mse_loc[0])):
    #         fig = plt.figure()
    #         mse_loc_mean = (np.asarray(mse_loc).mean(axis=2))
    #         timeslot = i*0.001
    #         plt.plot(delta_T, mse_loc_mean[:,i])
    #         plt.show()
if (args.animation):
    if args.sim_type == 'springcharge':
        sim.sample_trajectory_animation(T=args.length, sample_freq=args.sample_freq)
    else:
        print('Animation only implemented for springcharge Model')


import json
import os
import time
from models import Downloader, User

SETTINGS_FILE = "settings.json"

# Create history file if it does not exist
if not os.path.exists("history.txt"):
    with open("history.txt", "w") as f:
        f.write("")

# Check if the JSON file exists
if os.path.exists(SETTINGS_FILE):
    # If the file exists, open it and load its contents
    with open(SETTINGS_FILE, "r") as f:
        data = json.load(f)
else:
    # If the file doesn't exist, create it with default contents
    data = {"domain": None}
    with open(SETTINGS_FILE, "w") as f:
        json.dump(data, f)

if data["domain"] is None:
    print("Uzaktan eğitim websitesi adresini giriniz: (https:// olmadan)")
else:
    print(f"Uzaktan eğitim websitesi adresini giriniz: ({data['domain']} için enter)")

domain = input()
while domain == "" and data["domain"] is None:
    print("Lütfen geçerli bir adres giriniz.")
    domain = input()

if data["domain"] is None and domain != "":
    data["domain"] = domain
    with open(SETTINGS_FILE, "w") as f:
        json.dump(data, f)
elif data["domain"] is not None and domain == "":
    domain = data["domain"]

print("Çerezleri giriniz:")
cookies = input()
while cookies == "":
    print("Çerezler boş bırakılamaz.")
    cookies = input()

print("Ders listesi alınıyor.")
user = User(domain, cookies)
courses = user.get_enrolled_courses()
downloader = Downloader()
selected_activities = []
choose_again = True
previous_downloads = set(line.strip() for line in open("history.txt", encoding="utf-8"))
while len(courses) > 0 and choose_again:
    i = 1
    for course in courses:
        print(str(i) + ". " + course.name + " (" + str(course.activity_count) + ")")
        i += 1
    print("Ders seçiniz: ")
    course_selection = int(input())
    if 0 < course_selection <= len(courses):
        course = courses[course_selection - 1]
        print(course.name + " içeriği yükleniyor.")
        course.activities = []
        course.fetch_activities()
        i = 1
        for activity in course.activities:
            try:
                if activity.id in previous_downloads:
                    print(
                        "%d) Hafta: %d Dosya Adı: %s (Daha önce indirilmiş.)"
                        % (i, activity.weeks[0], activity.slug_name)
                    )
                else:
                    print(
                        "%d) Hafta: %d Dosya Adı: %s"
                        % (i, activity.weeks[0], activity.slug_name)
                    )
            except:
                continue
            i += 1
        print("İçerik aralığı seçiniz (Örn: 1-17): ")
        activity_selection = input().split("-")
        while len(activity_selection) != 2:
            print("Lütfen geçerli bir aralık giriniz.")
            activity_selection = input().split("-")
        if len(activity_selection) == 2:
            start_activity = int(activity_selection[0])
            end_activity = int(activity_selection[1])
            if (
                end_activity - start_activity >= 0
                and start_activity > 0
                and 0 < end_activity <= len(course.activities)
            ):
                for i in range(start_activity - 1, end_activity):
                    activity = course.activities[i]
                    activity.prepare_video(downloader)
                    selected_activities.append(activity)
                print("Başka bir dersten içerik seçmek ister misiniz? (e/h)")
                choose_again = input().lower() == "e"

if selected_activities:
    start = time.time()
    downloader.start_downloads()
    end = time.time()
else:
    print(selected_activities)
    print("İçerik seçilmedi.")


'''

This script contains functions which allow a user to select a ROI on a
single/multiple samples

These functions in a variety of ways call the roiselector function which is a GUI
to allow a user to select points. The functionionality of each funcitons is as follows:

    ChangePoint: allows a user to CHANGE the locations of found points corresponding
    to the matched images and their .feat files

    SelectArea: allows a user to propogate a selected area through a stack of images
    (works with any stack but probs best to use it after aligning tissue)

    SelectPoint: allows a user to SELECT a single point. used to manually annotate
    images during the featfinding period

'''

from time import sleep
import matplotlib
matplotlib.use('TkAgg')
import cv2
import numpy as np
# import tensorflow as tf
from glob import glob
import matplotlib.pyplot as plt
from matplotlib.widgets import RectangleSelector
import tifffile as tifi
from multiprocessing import Pool, cpu_count
from copy import deepcopy
from itertools import repeat
if __name__.find("HelperFunctions") == -1:
    from Utilities import *
else:
    from HelperFunctions.Utilities import *

# for each fitted pair, create an object storing their key information
class feature:
    def __init__(self, refP = None, tarP = None, dist = None, size = None, res = None, ID = None):
        # the position of the match on the reference image
        self.refP = refP

        # the position of the match on the target image
        self.tarP = tarP

        # eucledian error of the difference in gradient fields
        self.dist = dist

        # the size of the feature
        self.size = size

        # the resolution index of the image that was processed
        self.res = res

        self.ID = ID

    def __repr__(self):
        return repr((self.dist, self.refP, self.tarP, self.size, self.res))

def featChangePoint(dataSource, ref, tar, featureInfo = None, nopts = 5, ts = 4, title = "Matching"):

    # this fuction brings up a gui which allows for a user to manually CHANGE
    # features on the images. This modifies the original .feat file
    # Inputs:   (dataSource), the directory of where to get the info
    #           (ref, tar), reference and target name of the samples to change
    #           (nopts), number of point to change the annotations of, defaults 5
    #           (ts), text size for plots
    # Outputs:  (matchRef, matchTar), updated ref and target features with new points added
    #           (), also saves the new ref and tar positions in the SAME location

    # if modifying files
    if dataSource is None:
        imgref = ref
        imgtar = tar

        matchRefO = {}
        matchTarO = {}

        # if there are not features found, break
        if len(featureInfo) == 0:
            pass

        # if the info is in a list of objects
        elif type(featureInfo[0]) is not dict:
            for n, f in enumerate(featureInfo):
                matchRefO[n] = f.refP
                matchTarO[n] = f.tarP

        else:
            # if dictionaries are passed in as a list
            matchRefO = featureInfo[0]
            matchTarO = featureInfo[1]

    # if modifying arrays
    else:

        # get the dirs of the info
        infodirs = dataSource + 'info/'

        # if doing it for the main images
        try:
            imgdirs = dataSource + 'maskedSamples/'

            # get the images, doesn't m
            imgrefdir = glob(imgdirs + ref + ".*")[0]
            imgtardir = glob(imgdirs + tar + ".*")[0]

            imgref = cv2.imread(imgrefdir)
            imgtar = cv2.imread(imgtardir)

        # if doing it for the segSections
        except:
            imgdirs = dataSource
            imgrefdir = glob(imgdirs + ref + ".*")[0]
            imgtardir = glob(imgdirs + tar + ".*")[0]

            imgref = cv2.imread(imgrefdir)
            imgtar = cv2.imread(imgtardir)

        matchRefdir = infodirs + ref + ".reffeat"
        matchTardir = infodirs + tar + ".tarfeat"

        matchRefO = txtToDict(matchRefdir, float)[0]
        matchTarO = txtToDict(matchTardir, float)[0]

    # automatically set the text size
    ts = imgref.shape[0]/1000

    _, commonFeat = uniqueKeys([matchRefO, matchTarO])

    # create the dictionary with ONLY the common features and their positions
    # for the number of features specified by nofts
    matchRef = {}
    matchTar = {}

    # select only up to the specified number of points to use
    if nopts < len(commonFeat):
        commonFeat = commonFeat[:nopts]

    # if there are less matched points than desired, add some fake ones
    # to move around later
    elif nopts > len(commonFeat):
        for i in range(nopts - len(commonFeat)):
            pt = "added_" + str(i)
            commonFeat.append(pt)
            matchRefO[pt] = np.array([0, 0])
            matchTarO[pt] = np.array([0, 0])

    for cf in commonFeat:
        matchRef[cf] = matchRefO[cf]
        matchTar[cf] = matchTarO[cf]

    # create a standard combined image
    imgs = [imgref, imgtar]

    # get the image dimensions
    imgshapes = []
    for img in imgs:
        imgshapes.append(np.array(img.shape))

    # create a max size field of all images
    xm, ym, cm = np.max(np.array(imgshapes), axis = 0)
    field = np.zeros((xm, ym, cm)).astype(np.uint8)

    # stack the images next to each other to create a combined image
    imgsStand = []
    for img in imgs:
        xr, yr, c = img.shape
        img_s = field.copy(); img_s[:xr, :yr, :] = img
        imgsStand.append(img_s)
    imgCombine = np.hstack(imgsStand)

    featChange = {}
    cv2.startWindowThread()
    for feat in commonFeat:

        # get the pair of features to change
        featChange['ref'] = matchRef[feat]
        featChange['tar'] = matchTar[feat]

        for featC in featChange:

            featID = featChange[featC]

            # draw on the points
            imgCombineA = annotateImg(imgCombine.copy(), [matchRef, matchTar], ts)

            if featC == 'tar':
                featID += np.array([ym, 0])

            # highlight the point which is being changed
            cv2.circle(imgCombineA, tuple(featID.astype(int)), int(ts*12), (0, 255, 0), int(ts*8))

            # get the x and y position from the feature
            y, x = roiselector(imgCombineA, title)

            # if the window is closed to skip selecting a feature, keep the feature
            if (x * y == 0).all():
                yme = featID[0]
                xme = featID[1]
            else:
                yme = np.mean(y)
                xme = np.mean(x)

            # append reference and target information to the original list
            if featC == 'ref':
                matchRef[feat] = np.array((yme, xme))
            elif featC == "tar":
                matchTar[feat] = np.array((yme, xme)) - np.array([ym, 0])

    # save the new manually added positions to the original location, REPLACING the
    # information

    # return the data in the same format as the input
    if dataSource is None:

        # if the input was an object, return an object
        featInfos = []
        if type(featureInfo[0]) is not dict:
            for f in matchRef:
                featInfos.append(feature(refP = matchRef[f], tarP = matchTar[f], \
                    dist = -1, size = 100, res = -1, ID = None))

        else:
            featInfos = [matchRef, matchTar]

    else:
        dictToTxt(matchRef, matchRefdir, shape = imgref.shape, fit = False)
        dictToTxt(matchTar, matchTardir, shape = imgtar.shape, fit = False)

        # if the input was a list of dictionaries, return a list of dictionaries
        featInfos = [matchRef, matchTar]

    return(featInfos)

def featSelectArea(datahome, size, feats = 1, sample = 0, normalise = False, prefix = "tif"):

    # this function brings up a gui which allows user to manually selection a
    # roi on the image. This extracts samples from the aligned tissues and saves them

    cpuCount = int(cpu_count() * 0.75)
    segSections = datahome + str(size) + "/segSections/"

    serialised = False

    for f in range(feats):
        dirMaker(segSections + "seg" + str(f) + "/")

    alignedSamples = datahome + str(size) + "/RealignedSamples/"

    # get all the samples to be processed
    samples = sorted(glob(alignedSamples + "*" + prefix))

    # get the image to be used as the reference
    if type(sample) == int:
        refpath = samples[sample]
    elif type(sample) == str:
        refpath = glob(alignedSamples + sample + "*.tif")[0]

    # load the image
    try:
        img = cv2.cvtColor(tifi.imread(refpath), cv2.COLOR_BGR2RGB)
    except:
        img = cv2.imread(refpath)

    # create a small scale version of the image for color normalisation
    if normalise: imgref = cv2.resize(img, (int(img.shape[1] * 0.1), int(img.shape[0] * 0.1)))
    else: imgref = None

    # extract n feats from the target samples
    x = {}
    y = {}
    for f in range(feats):
        x[f], y[f] = roiselector(img)
        cv2.rectangle(img, (int(x[f][0]), int(y[f][0])), (int(x[f][1]), int(y[f][1])), (255, 255, 255), 40)
        cv2.rectangle(img, (int(x[f][0]), int(y[f][0])), (int(x[f][1]), int(y[f][1])), (0, 0, 0), 20)

    shapes = {}
    if serialised:
        for s in samples:
            name = nameFromPath(s)
            shapes[name] = sectionExtract(s, segSections, feats, x, y, imgref)

    else:
        with Pool(processes=cpuCount) as pool:
            shapes = pool.starmap(sectionExtract, zip(samples, repeat(segSections), repeat(feats), repeat(x), repeat(y), repeat(imgref)))

    # create a dictionary of all the tif shapes. they're all the same size,
    # its just about ensuring the input into align is consistent
    for i in range(feats):
        imgShapes = {}
        for n, s in enumerate(samples):
            name = nameFromPath(s, 3)
            imgShapes[name] = shapes[n][i]

        dictToTxt(imgShapes, segSections + "seg" + str(i) + "/info/all.tifshape")

def sectionExtract(path, segSections, feats, x, y, ref = None):

    img = cv2.imread(path)
    name = nameFromPath(path, 3)
    sections = []

    segShapes = {}
    # if a reference image is being used to normalise the images
    if (type(ref) is list) or (type(ref) is np.ndarray):
        img = hist_match(img, ref)

    for f in range(feats):

        print(name + " section " + str(f))
        segdir = segSections + "seg" + str(f) + "/"
        section = img[int(y[f][0]):int(y[f][1]), int(x[f][0]):int(x[f][1]), :]
        cv2.imwrite(segdir + name + ".tif", section)
        segShapes[f] = section.shape

    return(segShapes)

def roiselector(img, title = "Matching"):

    # function which calls the gui and get the co-ordinates
    # Inputs    (img), numpy array image of interest
    # Outputs   (xme, yme), x and y positions on the original image of the points selected
    # ensure that the image is scaled to a resolution which can be used in the sceen

    xc, yc, c = img.shape
    r = xc/yc
    # if the height is larger than the width
    if xc > yc/2:
        size = 700
        sizeY = int(size / r)
        sizeX = int(size)

    # if the width is larger than the height
    else:
        size = 1200
        sizeY = int(size)
        sizeX = int(size * r)

    scale = yc / sizeY

    # perform a search over a reduced size area
    imgr = cv2.cvtColor(cv2.resize(img, (sizeY, sizeX)), cv2.COLOR_RGB2BGR)
    x, y = annotatorGUI(imgr, title + " ENSURE YOU WAIT ~1 SEC BEFORE PRESSING ENTER")

    # get the postions
    # y = np.array([roi[1], roi[1] + roi[3]])
    # x = np.array([roi[0], roi[0] + roi[2]])

    # scale the positions back to their original size
    y = y * scale
    x = x * scale

    return(x, y)

def annotatorGUI(img, title):

    def line_select_callback(eclick, erelease):
        """
        Callback for line selection.

        *eclick* and *erelease* are the press and release events.
        """
        # get the click positions
        x1, y1 = eclick.xdata, eclick.ydata
        x2, y2 = erelease.xdata, erelease.ydata

        print(str([np.array([x1, x2]), np.array([y1, y2])]))

        # turn selection into point
        global allPos
        allPos = [np.array([x1, x2]), np.array([y1, y2])]

    def toggle_selector(event):
        if event.key == 'enter':
            toggle_selector.RS.set_active(False)
            fig.canvas.mpl_disconnect(toggle_selector)
            plt.close()
            # toggle_selector.RS.set_active(True)

    fig, ax = plt.subplots(figsize = (16, 12))
    global allPos
    allPos = [np.array([0, 0]), np.array([0, 0])]
    plt.axis("off")
    plt.imshow(img)
    ax.set_title(title)

    # drawtype is 'box' or 'line' or 'none'
    toggle_selector.RS = RectangleSelector(ax, line_select_callback,
                                        drawtype='box', useblit=True,
                                        button=[1, 3],  # disable middle button
                                        minspanx=1, minspany=1,
                                        spancoords='pixels',
                                        interactive=True)

    # get the positions by creating a global variable to access the callback function
    fig.canvas.mpl_connect('key_press_event', toggle_selector)
    plt.show()

    try:
        return(allPos)
    # if no new position declared, return the origin
    except:
        return([0, 0])

def annotateImg(imgs, info, ts):

    # this function takes an image and from a list of .feat dictionaries, draws
    # the position and information on and returns the images combined
    # Inputs:   (imgs), image type. if it is a list combine but if it is a single image just use as is
    #           (info), list of .feat dictionaries to annotate onto the images
    #           (ts), text size
    # Outputs:  (imgcombine), the combined image which has all the annotated features

    # if multiple images are being fed in then combine in a uniform array
    if type(imgs) is list:
        # get the image dimensions
        imgshapes = []
        for img in imgs:
            imgshapes.append(np.array(img.shape))

        # create a max size field of all images
        xm, ym, cm = np.max(np.array(imgshapes), axis = 0)
        field = np.zeros((xm, ym, cm)).astype(np.uint8)

        # stack the images next to each other to create a combined image
        imgsStand = []
        for img in imgs:
            xr, yr, c = img.shape
            img_s = field.copy(); img_s[:xr, :yr, :] = img
            imgsStand.append(img_s)
        imgCombine = np.hstack(imgsStand)

    # if a single image is being used then do all the annotations on this unmodified
    else:
        imgCombine = imgs
        xm, ym, cm = (np.array(imgCombine.shape) / len(info)).astype(int)

    for m, match in enumerate(info):

        for pos, r in enumerate(match):

            if type(match) is dict:
                r = match[r]

            # enusre that the co-ordinate is in the right format and position
            if type(r) is np.ndarray: r = tuple((r.astype(int)) + np.array([int(ym * m), 0]))
            else: r = tuple(np.array(r).astype(int) + np.array([int(ym * m), 0]))
            # if type(t) is np.ndarray: t = tuple(t.astype(int)

            # add the found points as marks
            cv2.circle(imgCombine, r, int(ts*10), (255, 0, 0), int(ts*5))

            # add point info to the image
            cv2.putText(imgCombine, str(pos),
                    tuple(r + np.array([20, 0])),
                    cv2.FONT_HERSHEY_SIMPLEX, ts, (255, 255, 255), int(ts*5))

            cv2.putText(imgCombine, str(pos),
                    tuple(r + np.array([20, 0])),
                    cv2.FONT_HERSHEY_SIMPLEX, ts, (0, 0, 0), int(ts*2.5))

    return(imgCombine)

if __name__ == "__main__":

    '''
    dataSource = '/Volumes/USB/H653/3/maskedSamples/'
    nameref = 'H653_01A_0.jpg'
    nametar = 'H653_02A_0.jpg'
    matchRef = {}
    matchTar = {}

    matchRef = {
    'feat_0':np.array([594, 347]),
    'feat_1':np.array([ 254, 1002]),
    'feat_2':np.array([ 527, 1163]),
    'feat_3':np.array([322, 262])
    }

    matchTar = {
    'feat_2':np.array([ 533, 1131]),
    'feat_3':np.array([294, 239]),
    'feat_4':np.array([287, 899]),
    'feat_5':np.array([608, 255])
    }

    featSelectPoint(nameref, nametar, matchRef, matchTar)
    '''

    dataSource = '/Volumes/USB/H653/'
    dataSource = '/Volumes/USB/H673A_7.6/'
    dataSource = '/Volumes/Storage/H653A_11.3new/'
    dataSource = '/Volumes/USB/H653A_11.3/'
    dataSource = '/Volumes/USB/H710C_6.1/'
    dataSource = '/Volumes/USB/H671A_18.5/'

    size = 3

    featChangePoint(r, f, ts = 4)
    # featSelectArea(dataSource, size, 2, 0, False)

    '''
    'H710C_289A+B_0',
    'H710C_304A+B_0',
    'H710C_308A+B_0',
    'H710C_308C_0',

    ref = [
        'H710C_311C_0'
    ]

    'H710C_289A+B_1',
    'H710C_304A+B_1',
    'H710C_308A+B_1',
    'H710C_309A+B_0',

    tar = [
        'H710C_312A+B_0'
    ]

    for r, f in zip(ref, tar):
        featChangePoint(r, f, ts = 4)
    '''