from DataStructures import Queue, Stack
from Position import Position
import pathfinder_gui

class Algorithms():

    def __init__(self):
        self.board = []
        self.startRow = 0
        self.startCol = 0
        self.finRow = 0
        self.finCol = 0
        self.numRow = 0
        self.numCol = 0
        self.finalPath = None

    def __setup(self):

        pBoard = pathfinder_gui.board
        self.numRow = len(pBoard)
        self.numCol = len(pBoard[0])
        self.board = [[Position(row, col) for col in range(self.numCol)]
                      for row in range(self.numRow)]
        self.startRow = 0
        self.startCol = 0
        self.finRow = 0
        self.finCol = 0
        self.finalPath = None

        for row in range(len(pBoard)):
            for col in range(len(pBoard[0])):
                if pBoard[row][col] == "S":
                    self.board[row][col].setType("START")
                    self.startRow = row
                    self.startCol = col
                elif pBoard[row][col] == "F":
                    self.board[row][col].setType("FINISH")
                    self.finRow = row
                    self.finCol = col
                elif pBoard[row][col] == "#":
                    self.board[row][col].setType("WALL")

    def Breadthfirst(self):
        self.__setup()
        queue = Queue()
        self.finalPath = None
        pathFound = False
        temp = None

        if self.board[self.startRow][self.startCol].getType() == "START":
            queue.enqueue(self.board[self.startRow][self.startCol])
            queue.getFront().getData().setVisited(True)

            while not queue.isEmpty() and not pathFound:
                temp = queue.dequeue().getData()

                if temp.getType() == "FINISH":
                    pathFound = True
                else:

                    if temp.getRow() > 0 and self.board[temp.getRow() - 1][temp.getCol()].getType() != "WALL"\
                            and self.board[temp.getRow() - 1][temp.getCol()].getVisited() == False:

                        self.board[temp.getRow() - 1][temp.getCol()
                                                      ].setVisited(True)
                        self.board[temp.getRow() - 1][temp.getCol()
                                                      ].setPrev(temp)
                        queue.enqueue(
                            self.board[temp.getRow() - 1][temp.getCol()])

                    if temp.getRow() < self.numRow - 1\
                            and self.board[temp.getRow() + 1][temp.getCol()].getType() != "WALL"\
                            and self.board[temp.getRow() + 1][temp.getCol()].getVisited() == False:

                        self.board[temp.getRow() + 1][temp.getCol()
                                                      ].setVisited(True)
                        self.board[temp.getRow() + 1][temp.getCol()
                                                      ].setPrev(temp)
                        queue.enqueue(
                            self.board[temp.getRow() + 1][temp.getCol()])

                    if temp.getCol() > 0 and self.board[temp.getRow()][temp.getCol() - 1].getType() != "WALL"\
                            and self.board[temp.getRow()][temp.getCol() - 1].getVisited() == False:

                        self.board[temp.getRow()][temp.getCol() -
                                                  1].setVisited(True)
                        self.board[temp.getRow()][temp.getCol() -
                                                  1].setPrev(temp)
                        queue.enqueue(
                            self.board[temp.getRow()][temp.getCol() - 1])

                    if temp.getCol() < self.numCol - 1\
                            and self.board[temp.getRow()][temp.getCol() + 1].getType() != "WALL"\
                            and self.board[temp.getRow()][temp.getCol() + 1].getVisited() == False:

                        self.board[temp.getRow()][temp.getCol() +
                                                  1].setVisited(True)
                        self.board[temp.getRow()][temp.getCol() +
                                                  1].setPrev(temp)
                        queue.enqueue(
                            self.board[temp.getRow()][temp.getCol() + 1])
                pathfinder_gui.showQueue(temp, queue)

        if pathFound:
            self.finalPath = Stack()

            while temp != None:
                self.finalPath.push(temp)
                temp = temp.getPrev()

            pathfinder_gui.showPath(self.finalPath)

    def Depthfirst(self):
        self.__setup()
        stack = Stack()
        temp = None
        pathFound = False
        self.finalPath = None

        if self.board[self.startRow][self.startCol].getType() == "START":
            stack.push(self.board[self.startRow][self.startCol])
            stack.getTop().getData().setVisited(True)

            while not stack.isEmpty() and not pathFound:
                temp = stack.pop().getData()

                if temp.getType() == "FINISH":
                    pathFound = True
                else:

                    if temp.getRow() > 0 and self.board[temp.getRow() - 1][temp.getCol()].getType() != "WALL"\
                            and self.board[temp.getRow() - 1][temp.getCol()].getVisited() == False:

                        self.board[temp.getRow() - 1][temp.getCol()
                                                      ].setVisited(True)
                        self.board[temp.getRow() - 1][temp.getCol()
                                                      ].setPrev(temp)
                        stack.push(
                            self.board[temp.getRow() - 1][temp.getCol()])

                    if temp.getRow() < self.numRow - 1\
                            and self.board[temp.getRow() + 1][temp.getCol()].getType() != "WALL"\
                            and self.board[temp.getRow() + 1][temp.getCol()].getVisited() == False:

                        self.board[temp.getRow() + 1][temp.getCol()
                                                      ].setVisited(True)
                        self.board[temp.getRow() + 1][temp.getCol()
                                                      ].setPrev(temp)
                        stack.push(
                            self.board[temp.getRow() + 1][temp.getCol()])

                    if temp.getCol() > 0 and self.board[temp.getRow()][temp.getCol() - 1].getType() != "WALL"\
                            and self.board[temp.getRow()][temp.getCol() - 1].getVisited() == False:

                        self.board[temp.getRow()][temp.getCol() -
                                                  1].setVisited(True)
                        self.board[temp.getRow()][temp.getCol() -
                                                  1].setPrev(temp)
                        stack.push(self.board[temp.getRow()]
                                   [temp.getCol() - 1])

                    if temp.getCol() < self.numCol - 1\
                            and self.board[temp.getRow()][temp.getCol() + 1].getType() != "WALL"\
                            and self.board[temp.getRow()][temp.getCol() + 1].getVisited() == False:

                        self.board[temp.getRow()][temp.getCol() +
                                                  1].setVisited(True)
                        self.board[temp.getRow()][temp.getCol() +
                                                  1].setPrev(temp)
                        stack.push(self.board[temp.getRow()]
                                   [temp.getCol() + 1])
                pathfinder_gui.showStack(temp, stack)

        if pathFound:
            self.finalPath = Stack()

            while temp != None:
                self.finalPath.push(temp)
                temp = temp.getPrev()
            pathfinder_gui.showPath(self.finalPath)

    def aStar(self):
        self.__setup()

        availNodes = []
        usedNodes = []
        availNodes.append(self.board[self.startRow][self.startCol])

        pathFound = False
        while not pathFound and len(availNodes) > 0:
            neighbours = []
            currNode = availNodes[0]
            availNodes.remove(currNode)
            usedNodes.append(currNode)

            if currNode.getType() == "FINISH":  # path has been found
                pathFound = True
            else:
                if currNode.getRow() > 0 and self.board[currNode.getRow() - 1][currNode.getCol()].getType() != "WALL":
                    neighbours.append(
                        self.board[currNode.getRow() - 1][currNode.getCol()])

                if currNode.getCol() > 0 and self.board[currNode.getRow()][currNode.getCol() - 1].getType() != "WALL":
                    neighbours.append(
                        self.board[currNode.getRow()][currNode.getCol() - 1])

                if currNode.getRow() < self.numRow - 1 and self.board[currNode.getRow() + 1][currNode.getCol()].getType() != "WALL":
                    neighbours.append(
                        self.board[currNode.getRow() + 1][currNode.getCol()])

                if currNode.getCol() < self.numCol - 1 and self.board[currNode.getRow()][currNode.getCol() + 1].getType() != "WALL":
                    neighbours.append(
                        self.board[currNode.getRow()][currNode.getCol() + 1])

                for neighbour in neighbours:
                    if neighbour not in usedNodes and (currNode.getGCost() + 10 < neighbour.getGCost() or neighbour not in availNodes):
                        neighbour.setGCost(currNode.getGCost() + 10)
                        neighbour.setHCost(int(((abs(neighbour.getRow() - self.finRow)**2 +
                                                 abs(neighbour.getCol() - self.finCol)**2)**(0.5))*10))

                        neighbour.setPrev(currNode)
                        if neighbour not in availNodes:
                            self.__sortedInsert(availNodes, neighbour)
            pathfinder_gui.showAStar(currNode, availNodes)

        if pathFound:
            self.finalPath = Stack()

            while currNode != None:
                self.finalPath.push(currNode)
                currNode = currNode.getPrev()

            pathfinder_gui.showPath(self.finalPath)

    def __sortedInsert(self, toAddto, elem):

        for a in range(len(toAddto)):
            if elem.getFCost() < toAddto[a].getFCost():
                toAddto.insert(a, elem)
                return
            elif elem.getFCost() == toAddto[a].getFCost():
                if elem.getHCost() < toAddto[a].getHCost():
                    toAddto.insert(a, elem)
                else:
                    toAddto.insert(a + 1, elem)
                return
        toAddto.append(elem)


import os
import re
import time
import logging
import hashlib
import json

import dataclasses
from dataclasses import dataclass, field as dfield
from collections import defaultdict
from threading import RLock, Thread, local
from typing import Callable, Dict, List, Optional

import idzip

import ms_deisotope
from ms_deisotope.data_source.scan.base import ScanBase
from ms_deisotope.data_source.scan.loader import RandomAccessScanSource

from ms_deisotope.utils import LRUDict
from ms_deisotope.data_source._compression import get_opener
from ms_deisotope.data_source.usi import USI

BLOCK_SIZE = 2 ** 20
logger = logging.getLogger(__name__)

@dataclass(frozen=True)
class FileType:
    accession: str
    name: str

    def __str__(self):
        return self.name

    def to_dict(self) -> Dict:
        return dataclasses.asdict(self)

    @classmethod
    def from_dict(cls, state):
        return cls(**state)

PeakListFile = FileType("MS:1002850", "Peak list file URI")
AssociatedRawFile = FileType("MS:1002846", "Associated raw file URI")

@dataclass(frozen=True)
class PROXIDataFile:
    uri: str
    file_type: FileType

    def to_dict(self) -> Dict:
        d = self.file_type.to_dict()
        d['value'] = self.uri
        return d

    @classmethod
    def from_dict(cls, d):
        ft = FileType(d['accession'], d['name'])
        return cls(d['value'], ft)

@dataclass
class PROXIDataset:
    accession: str = dfield(default=None)
    title: str = dfield(default=None)
    summary: str = dfield(default=None)
    instruments: List[str] = dfield(default_factory=list)
    identifiers: List[str] = dfield(default_factory=list)
    species: List[str] = dfield(default_factory=list)
    modifications: List[str] = dfield(default_factory=list)
    contacts: List[str] = dfield(default_factory=list)
    publications: List[str] = dfield(default_factory=list)
    keywords: List[str] = dfield(default_factory=list)
    dataset_link: str = dfield(default=None)
    dataset_files: List[PROXIDataFile] = dfield(default_factory=list)
    links: List[str] = dfield(default_factory=list)

    def to_dict(self):
        d = dataclasses.asdict(self)
        d['datasetLinks'] = d.pop("dataset_links", None)
        # TODO: Make this consistent whether or not it is a dict or an instance
        d['dataFiles'] = [PROXIDataFile(f['uri'], FileType.from_dict(f['file_type'])).to_dict() for f in d.pop('dataset_files', [])]
        return d

    @classmethod
    def from_dict(cls, d):
        d = d.copy()
        d['dataset_links'] = d.pop('datasetLinks', None)
        d['dataset_files'] = [
            PROXIDataFile.from_dict(f) for f in d.pop('datasetFiles', [])]
        return cls(**d)

    @classmethod
    def new(cls, dataset_id):
        self = cls()
        self.accession = dataset_id
        self.title = dataset_id
        self.summary = dataset_id
        return self

class PROXIDatasetAPIMixinBase(object):

    def _datafile_uri(self, dataset_id, data_file):
        raise NotImplementedError()

    def _is_peaklist(self, d):
        base = os.path.basename(d)
        if base.endswith(".gz"):
            base = base[:-3]
        base = base.lower()
        _, ext = base.rsplit(".", 1)
        if ext in tuple(map(str.lower, self.peaklist_file_types)):
            return True
        return False

    @property
    def peaklist_file_types(self):
        """
        Alias to allow :attr:`valid_file_types` to vary.

        Returns
        -------
        list[str]
        """
        return self.valid_file_types

    def describe_dataset(self, dataset_id: str) -> PROXIDataset:
        record = PROXIDataset.new(dataset_id)
        data_files = self._find_mz_files(self.index[dataset_id])
        for d in data_files:
            uri = self._datafile_uri(dataset_id, d)
            if self._is_peaklist(d):
                f = PROXIDataFile(uri, PeakListFile)
            else:
                f = PROXIDataFile(uri, AssociatedRawFile)
            record.dataset_files.append(f)
        return record

    def enumerate_datasets(self, result_type=None):
        if result_type == 'compact':
            result = list(self.index.keys())
            return result
        return [
            self.describe_dataset(d).to_dict() for d in self.index.keys()]

class IndexManagingMixin(object):

    def __init__(self, **kwargs):
        self.index = self.build_index()
        self.monitor_thread = self.monitor_index(
            kwargs.get("update_frequency"))

    def _monitor_index(self, update_frequency=None):
        if update_frequency is None:
            # Default to one hour, 60^2 seconds
            update_frequency = 60.0 ** 2
        logger.info(
            "Monitoring archive for changes every %0.2f seconds.", update_frequency)
        while True:
            time.sleep(update_frequency)
            logger.info("Rebuilding Index!")
            self.index = self.build_index(use_cache=False)

    def monitor_index(self, update_frequency=None):
        update_thread = Thread(
            target=self._monitor_index,
            args=(update_frequency, ))
        update_thread.daemon = True
        update_thread.start()
        return update_thread

    def _checksum(self, string):
        return hashlib.new("md5", string).hexdigest()

    def _find_index_cache(self, use_cache=True):
        cache_name = hashlib.md5(
            self.base_uri.encode('utf8')).hexdigest() + '.json'
        if os.path.exists(cache_name) and use_cache:
            logger.info("Reading cache from %r", cache_name)
            try:
                with open(cache_name, 'rt') as fh:
                    index = defaultdict(list, json.load(fh))
                    return cache_name, index
            except IOError:
                logger.error("Failed to read index from disk", exc_info=True)
        return cache_name, None

class SpectrumSerializerServiceMixin(object):
    scan_filters: List[Callable[[
        ScanBase, Optional[RandomAccessScanSource]], ScanBase]]

    def __init__(self, scan_filters: Optional[List]=None):
        self.scan_filters = scan_filters or []

    def process_scan(self, scan: ScanBase, reader: Optional[RandomAccessScanSource]):
        """
        Process ``scan`` with each callable in :attr:`scan_filters` prior to serialization.

        Arguments
        ---------
        scan : :class:`~.ScanBase`
            The scan being processed
        reader : :class:`~.RandomAccessScanSource`
            The data source for ``scan``

        Returns
        -------
        :class:`~.ScanBase`
        """
        for filt in self.scan_filters:
            scan = filt(scan, reader)
        return scan

    def add_scan_filter(self, filt: Callable[[ScanBase, Optional[RandomAccessScanSource]], ScanBase]):
        """
        Add a new scan filtering callable to the chain of filters.

        This new filter is applied after all previous filters.

        Parameters
        ----------
        filt : Callable[[ScanBase, Optional[RandomAccessScanSource]], ScanBase]
            A callable object that transforms a :class:`~.ScanBase` object with optional access to its
            source.
        """
        self.scan_filters.append(filt)

    def build_response(self, usi: USI, scan: ScanBase):
        """
        Build a response :class:`dict` representing ``scan`` for the requested ``usi``.

        Parameters
        ----------
        usi : :class:`~pyteomics.usi.USI`
            The parsed USI requested.
        scan : :class:`~.ScanBase`
            The scan retrieved for this request.

        Returns
        -------
        dict
        """
        payload = {
            "usi": str(usi),
            "status": "READABLE",
            "attributes": [
                {"accession": "MS:1000511", "name": "ms level",
                    "value": scan.ms_level},
                {"accession": "MS:1008025", "name": "scan number",
                    "value": scan.index + 1},
                {"accession": "MS:1003061", "name": "spectrum name", "value": scan.id},
            ]
        }
        payload = self._write_peak_arrays(usi, scan, payload)
        if scan.ms_level > 1:
            mz = None
            charge = None
            if scan.precursor_information:
                mz = scan.precursor_information.mz
                charge = scan.precursor_information.charge
            elif scan.isolation_window:
                mz = scan.isolation_window.target
            if mz:
                payload['attributes'].append({
                    "accession": "MS:1000827", "name": "isolation window target m/z",
                    "value": mz,
                })
            if charge:
                payload['attributes'].append({
                    "accession": "MS:1000041", "name": "charge state",
                    "value": charge,
                })
        return payload

    def _write_peak_arrays(self, usi, scan, payload):
        mzs = []
        intensities = []
        charges = None
        base_peak = None
        nb_peaks = None
        if scan.deconvoluted_peak_set is not None:
            mzs = [p.mz for p in scan.deconvoluted_peak_set]
            intensities = [p.intensity for p in scan.deconvoluted_peak_set]
            charges = [p.charge for p in scan.deconvoluted_peak_set]
            is_profile = False
            nb_peaks = len(scan.deconvoluted_peak_set)
            base_peak = scan.base_peak.deconvoluted().intensity
        elif scan.peak_set is not None:
            mzs = [p.mz for p in scan.peak_set]
            intensities = [p.intensity for p in scan.peak_set]
            charges = None
            is_profile = False
            nb_peaks = len(scan.peak_set)
            base_peak = scan.base_peak.centroided().intensity
        else:
            mzs = scan.arrays.mz.tolist()
            intensities = scan.arrays.intensity.tolist()
            is_profile = scan.is_profile
            nb_peaks = len(mzs)
            base_peak = scan.base_peak.raw().intensity

        payload['m/z array'] = mzs
        payload['intensity array'] = intensities
        if charges is not None:
            payload['charge array'] = charges
        payload['attributes'].append(
            {"accession": "MS:1000128", "name": "profile spectrum"} if is_profile else
            {"accession": "MS:1000127", "name": "centroid spectrum"})
        payload['attributes'].append({
            "accession": 'MS:1000505', 'name': 'base peak intensity', 'value': base_peak
        })
        payload['attributes'].append({
            "accession": 'MS:1003059', 'name': 'number of peaks', 'value': nb_peaks
        })
        return payload

class DatasetNotAvailable(ValueError):
    """An error code indicating that the requested dataset isn't available."""

class UnrecognizedIndexFlag(ValueError):
    """An error code indicating that the USI index flag was not recognized."""

class MassSpectraDataArchiveBase(SpectrumSerializerServiceMixin, PROXIDatasetAPIMixinBase, IndexManagingMixin):
    valid_file_types = ['mzML', 'mzXML', 'mgf', 'mzMLb']
    special_openers = set()

    def __init__(self, base_uri, cache_size=None, scan_filters: Optional[List] = None, **kwargs):
        if cache_size is None:
            cache_size = 1
        self.base_uri = base_uri
        self.cache = LRUDict(cache_size=cache_size)
        IndexManagingMixin.__init__(self, **kwargs)
        SpectrumSerializerServiceMixin.__init__(self, scan_filters=scan_filters)

    def build_index(self, use_cache=True):
        file_types = self.valid_file_types + ['json']

        cache_name, index = self._find_index_cache(use_cache=use_cache)
        if index is not None and use_cache:
            return index

        index = defaultdict(list)

        self._traverse_files(file_types, index)

        reduced = {
            k: self._find_mz_files(v) for k, v in index.items()
        }

        n_datasets = len(reduced)
        n_ms_files = 0
        for _, v in reduced.items():
            n_ms_files += len(v)

        logger.info(
            "Collected %d Datasets with %d MS Data Files in Index", n_datasets, n_ms_files)
        logger.info("Saving index to %r", cache_name)
        try:
            with open(cache_name, 'wt') as fh:
                json.dump(index, fh, indent=2, sort_keys=True)
        except IOError:
            logger.error("Failed to write index to disk", exc_info=True)
        return index

    def _traverse_files(self, file_types, index):
        raise NotImplementedError()

    def find_ms_files_for(self, dataset, run_name):
        result = []
        members = self.index.get(dataset)
        if run_name is None:
            return members

        if members is None:
            raise DatasetNotAvailable((dataset, run_name))

        for member in members:
            base = member.split("/")[-1]
            if base.startswith(run_name):
                result.append(member)
            elif base.endswith("json"):
                if base.startswith(run_name.rsplit(".", 1)[0]):
                    result.append(member)
        # Validate that there is only one MS file for the name
        i = 0
        for res in sorted(result):
            if re.search(r"-byte-offsets.json$", res):
                continue
            if re.search(r"-idx.json$", res):
                continue
            i += 1
        if i > 1:
            raise ValueError("Multiple MS Data Files Found")
        return self._find_mz_files(result, first=True), self._find_index_files(result, first=True)

    def _find_mz_files(self, namelist, first=False):
        pattern = re.compile(r"(?:%s)(?:.gz)?$" %
                             "|".join(self.valid_file_types))
        names = []
        for name in namelist:
            if pattern.search(name):
                if first:
                    return name
                names.append(name)
        if first and not names:
            return None
        return names

    def _find_index_files(self, namelist, first=False):
        names = []
        for name in namelist:
            if re.search(r"-byte-offsets.json$", name):
                if first:
                    return name
                names.append(name)
        if first and not names:
            return None
        return names

    def _opener(self, uri, block_size=None, **kwargs):
        raise NotImplementedError()

    def open_ms_file(self, ms_file_uri, index_uri=None):
        if ms_file_uri in self.cache:
            logger.info("Found %r in the cache", ms_file_uri)
            return self.cache[ms_file_uri]
        logger.info("Opening %r", ms_file_uri)
        base_name =  ms_file_uri.split("/")[-1]
        if base_name.split('.')[-1] in self.special_openers:
            reader = ms_deisotope.MSFileLoader(ms_file_uri)
        else:
            mzml_fh = self._opener(ms_file_uri)
            if ms_file_uri.endswith("gz"):
                mzml_fh = idzip.IdzipFile(fileobj=mzml_fh)
            index_fh = None
            if index_uri is not None:
                index_fh = self._opener(index_uri, block_size=BLOCK_SIZE)
            else:
                logger.warning("Byte offset index is missing for %r", ms_file_uri)
            reader = ms_deisotope.MSFileLoader(mzml_fh, index_file=index_fh)
        logger.info("Finished opening %r", ms_file_uri)
        lock = RLock()
        self.cache[ms_file_uri] = (reader, lock)
        return reader, lock

    def list_collections(self):
        return list(self.index.keys())

    def list_datasets(self):
        return {
            k: [f.split("/")[-1] for f in self._find_mz_files(v)] for k, v in self.index.items()
        }

    def get_scan_from_usi(self, usi, reader):
        if usi.scan_identifier_type == 'scan':
            index = int(usi.scan_identifier) - 1
        elif usi.scan_identifier_type == "index":
            index = int(usi.scan_identifier)
        else:
            raise UnrecognizedIndexFlag(usi.scan_identifier_type)
        scan = reader.get_scan_by_index(index)
        return scan

    def enumerate_spectra_for(self, dataset, ms_run=None):
        raise NotImplementedError()

    def handle_usi(self, usi, convert_json=True):
        """
        Receive a USI, parse it into its constituent pieces, locate
        the relevant spectrum, and return it formatted message.

        Parameters
        ----------
        usi : str
            The identifier of the spectrum to be retrieved
        convert_json : bool, optional
            Whether to format the spectrum as JSON or not. Defaults to :const:`True`.

        Returns
        -------
        spectrum : dict or Scan
            The spectrum, rendered according to the `convert_json` or not.
        """
        usi = USI.parse(str(usi))
        mzml_uri, index_uri = self.find_ms_files_for(usi.dataset, usi.datafile)
        reader, lock = self.open_ms_file(mzml_uri, index_uri)
        with lock:
            scan = self.get_scan_from_usi(usi, reader)
            scan = self.process_scan(scan, reader)
            if convert_json:
                payload = self.build_response(usi, scan)
                return payload
            else:
                return scan

    def __call__(self, usi):
        return self.handle_usi(usi)

class FSMassSpectraDataArchive(MassSpectraDataArchiveBase):
    def __init__(self, base_uri, cache_size=None, **kwargs):
        self.valid_file_types = self.valid_file_types[:]
        from ms_deisotope.data_source.thermo_raw_net import determine_if_available
        if determine_if_available():
            self.valid_file_types.append("raw")
            self.special_openers = set(self.special_openers) | {"raw"}
        super(FSMassSpectraDataArchive, self).__init__(
            base_uri, cache_size=cache_size, **kwargs)

    def _opener(self, uri, block_size=None, **kwargs):
        return open(uri, 'rb')

    def _datafile_uri(self, dataset_id, data_file):
        return os.path.normpath('/'.join([self.base_uri, data_file])).replace(os.sep, '/')

    def _traverse_files(self, file_types, index):
        for prefix, dirs, files in os.walk(self.base_uri):
            if os.sep in prefix:
                tokens = prefix.split(os.sep)
            else:
                tokens = prefix.split("/")
            if tokens[-1] == '':
                base = tokens[-2]
            else:
                base = tokens[-1]

            for f in files:
                try:
                    ext = f.rsplit(".", 1)[1]
                except IndexError:
                    continue
                if ext == 'gz':
                    ext = f.rsplit(".", 2)[1]
                if ext in file_types:
                    index[base].append(
                        os.path.normpath(
                            os.path.join(prefix, f)).replace(
                                os.sep, "/"))

class S3MassSpectraDataArchive(MassSpectraDataArchiveBase):
    def __init__(self, base_uri, cache_size=None, s3_args=None, **kwargs):
        if s3_args is None:
            s3_args = {}
        self._tls = local()
        self.s3_args = s3_args
        super(S3MassSpectraDataArchive, self).__init__(
            base_uri=base_uri, cache_size=cache_size, **kwargs)

    @property
    def s3_service(self):
        try:
            return self._tls.s3_service
        except (AttributeError, KeyError):
            from s3fs import S3FileSystem
            self._tls.s3_service = S3FileSystem(**self.s3_args)
            return self._tls.s3_service

    def _opener(self, uri, block_size=None, **kwargs):
        if block_size is None:
            block_size = BLOCK_SIZE
        return self.s3_service.open(uri, block_size=block_size)

    def _datafile_uri(self, dataset_id, data_file):
        return "s3://" + data_file

    def _traverse_files(self, file_types, index):
        for prefix, _dirs, files in self.s3_service.walk(self.base_uri):
            for f in files:
                query = f
                if query.endswith(".gz"):
                    query = query[:-3]
                ext = query.rsplit(".", 1)[-1]
                for t in file_types:
                    if t == ext:
                        index[prefix.replace(self.base_uri[5:], '')].append(
                            prefix + '/' + f)
                        break
        for group, members in list(index.items()):
            if "/" in group:
                base, _rest = group.split("/", 1)
                index[base].extend(members)
                index.pop(group)

class PeakPickingScanFilter(object):
    def __init__(self, peak_picking_args=None, *args, **kwargs):
        if peak_picking_args is None:
            peak_picking_args = {}
        self.peak_picking_args = peak_picking_args

    def __call__(self, scan, reader):
        scan.pick_peaks(**self.peak_picking_args)
        return scan

class DeconvolutingScanFilter(PeakPickingScanFilter):
    def __init__(self, *args, averagine=ms_deisotope.peptide, scorer=ms_deisotope.MSDeconVFitter(0), truncate_after: float=0.8, **kwargs):
        self.deconvolution_args = kwargs.get('deconvolution_args', {})
        self.deconvolution_args.setdefault("averagine", averagine)
        self.deconvolution_args.setdefault("scorer", scorer)
        self.deconvolution_args.setdefault("truncate_after", truncate_after)
        super(DeconvolutingScanFilter, self).__init__(*args, **kwargs)

    def __call__(self, scan, reader):
        scan = super().__call__(scan, reader)
        scan.deconvolute(**self.deconvolution_args)
        return scan


import cv2
import imutils

# lower = (165, 126, 211)
# upper = (185, 146, 291)

# 152 181 151

# grab the current frame
frame = cv2.imread('dial.jpg')
# handle the frame from VideoCapture or VideoStream

frame_lab = cv2.cvtColor(frame,cv2.COLOR_BGR2LAB)
frame_hsv = cv2.cvtColor(frame,cv2.COLOR_BGR2HSV)

# ROIs = cv2.selectROIs('Select ROIs',frame, False)
# frame = frame[ROIs[0][1]:ROIs[0][1]+ROIs[0][3], ROIs[0][0]:ROIs[0][0]+ROIs[0][2]]

# resize the frame, blur it, and convert it to the HSV
# color space
# print(ROI_1)
# cv2.imshow('frae',frame)
# cv2.imshow('roi',ROI_1)

import numpy as np

# color = (10,181,151)

# hsv = cv2.cvtColor(blurred, cv2.COLOR_BGR2HSV)

lower_lab = (10,181,151)
upper_lab = (152,181,151)
# print(frame.mean())

lower_bgr = cv2.cvtColor( np.uint8([[lower_lab]] ), cv2.COLOR_LAB2BGR)[0][0]
lower_hsv = cv2.cvtColor( np.uint8([[lower_bgr]] ), cv2.COLOR_BGR2HSV)[0][0]

upper_bgr = cv2.cvtColor( np.uint8([[upper_lab]] ), cv2.COLOR_LAB2BGR)[0][0]
upper_hsv = cv2.cvtColor( np.uint8([[upper_bgr]] ), cv2.COLOR_BGR2HSV)[0][0]

print(lower_hsv, upper_hsv)
# img = np.full((100,100,3),color, dtype=np.uint8)
# img = cv2.cvtColor(img, cv2.COLOR_LAB2BGR)
# print(img)
# cv2.imshow('np',img)

mask = cv2.inRange(frame_hsv, lower_hsv, upper_hsv)
mask = cv2.erode(mask, None, iterations=2)
mask = cv2.dilate(mask, None, iterations=2)
cv2.imshow('mask', mask)
cv2.imshow('hsv', frame_hsv)

# 152 181 151

cv2.waitKey(0)


import os
import numpy as np
from pathlib import Path
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.init
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms

# set parameters
epochs = 10
lr = 1e-3
batch_size = 32
chkpt_dir = "./checkpoint"
log_dir = "./log"
device = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(777)
if device == "cuda":
    torch.cuda.manual_seed_all(777)

# define network

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=32,
                      kernel_size=3, padding=1, bias=True),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(in_channels=32, out_channels=64,
                      kernel_size=3, padding=1, bias=True),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.fc = nn.Linear(in_features=3136, out_features=10, bias=True)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = x.view(x.shape[0], -1)
        x = self.fc(x)
        return x

# save and load network

def save(chkpt_dir: str, net: Net, optim: optim.Adam, epoch: int):
    if not os.path.exists(chkpt_dir):
        os.makedirs(chkpt_dir)

    net_dict = net.state_dict()
    optim_dict = optim.state_dict()
    model_dict = dict(net=net_dict, optim=optim_dict)
    save_path = f"{chkpt_dir}/model_epoch_{epoch}.pth"
    torch.save(model_dict, save_path)

def load(chkpt_dir: str, net: Net, optim: optim.Adam):
    chkpt_list = os.listdir(chkpt_dir)
    filename_lst = [os.path.splitext(path)[0] for path in chkpt_list]
    idxs = [filename.split("_")[-1] for filename in filename_lst]
    last_idx = np.argmax(idxs)
    model_dict = torch.load(f"{chkpt_dir}/{chkpt_list[last_idx]}")

    net.load_state_dict(model_dict["net"])
    optim.load_state_dict(model_dict["optim"])

# get datasets
mnist_train = datasets.MNIST(
    root="./",
    download=True,
    train=True,
    transform=transforms.ToTensor()
)
loader = DataLoader(
    mnist_train,
    batch_size=batch_size,
    shuffle=True,
    drop_last=True,
)

# set network, optimizer, tensorboard writer
net = Net().to(device)
criterion = nn.CrossEntropyLoss().to(device)
optimizer = optim.Adam(params=net.parameters(), lr=lr)
writer = SummaryWriter(log_dir=log_dir)

# train
for epoch in range(1, epochs+1):
    loss_arr = []
    for batch, (input, label) in enumerate(loader):
        input.to(device)
        label.to(device)

        optimizer.zero_grad()
        pred = net(input)
        loss = criterion(pred, label)
        loss.backward()
        optimizer.step()

        loss_arr.append(loss.item())
        print(
            f"train: EPOCH {epoch}/{epochs} | BATCH {batch}/{len(loader)} | Loss {np.mean(loss_arr):.4f}")

    writer.add_scalar("loss", np.mean(loss_arr), epoch)

writer.close()


import torch

from ..builder import BBOX_ASSIGNERS, build_bbox_coder
from ..iou_calculators import build_iou_calculator
from .assign_result import AssignResult
from .base_assigner import BaseAssigner

@BBOX_ASSIGNERS.register_module()
class FPNRankingAssigner(BaseAssigner):
    """Assign a corresponding gt bbox or background to each bbox.

    Each proposals will be assigned with `-1`, or a semi-positive integer
    indicating the ground truth index.

    - -1: negative sample, no assigned gt
    - semi-positive integer: positive sample, index (0-based) of assigned gt

    Args:
        pos_iou_thr (float): IoU threshold for positive bboxes.
        neg_iou_thr (float or tuple): IoU threshold for negative bboxes.
        min_pos_iou (float): Minimum iou for a bbox to be considered as a
            positive bbox. Positive samples can have smaller IoU than
            pos_iou_thr due to the 4th step (assign max IoU sample to each gt).
        gt_max_assign_all (bool): Whether to assign all bboxes with the same
            highest overlap with some gt to that gt.
        ignore_iof_thr (float): IoF threshold for ignoring bboxes (if
            `gt_bboxes_ignore` is specified). Negative values mean not
            ignoring any bboxes.
        ignore_wrt_candidates (bool): Whether to compute the iof between
            `bboxes` and `gt_bboxes_ignore`, or the contrary.
        match_low_quality (bool): Whether to allow low quality matches. This is
            usually allowed for RPN and single stage detectors, but not allowed
            in the second stage. Details are demonstrated in Step 4.
        gpu_assign_thr (int): The upper bound of the number of GT for GPU
            assign. When the number of gt is above this threshold, will assign
            on CPU device. Negative values mean not assign on CPU.
    """

    def __init__(self,
                 gt_max_assign_all=True,
                 ignore_iof_thr=-1,
                 ignore_wrt_candidates=True,
                 gpu_assign_thr=512,
                 iou_calculator=dict(type='BboxOverlaps2D'),
                 assign_metric='iou',
                 topk=1,
                 topq=10,
                 cls_w = 1.0,
                 inside_circle=False,
                 gauss_thr = 1.5,
                 bbox_coder=dict(
                     type='DeltaXYWHAOBBoxCoder',
                     target_means=(.0, .0, .0, .0, .0),
                     target_stds=(1.0, 1.0, 1.0, 1.0, 1.0))):
        self.gt_max_assign_all = gt_max_assign_all
        self.ignore_iof_thr = ignore_iof_thr
        self.ignore_wrt_candidates = ignore_wrt_candidates
        self.gpu_assign_thr = gpu_assign_thr
        self.iou_calculator = build_iou_calculator(iou_calculator)
        self.assign_metric = assign_metric
        self.topk = topk
        self.topq = topq
        self.cls_w = cls_w
        self.inside_circle = inside_circle
        self.gauss_thr = gauss_thr
        self.bbox_coder = build_bbox_coder(bbox_coder)

    def assign(self, cls_scores, bbox_preds, bboxes, gt_bboxes, gt_bboxes_ignore=None, gt_labels=None):
        """Assign gt to bboxes.

        This method assign a gt bbox to every bbox (proposal/anchor), each bbox
        will be assigned with -1, or a semi-positive number. -1 means negative
        sample, semi-positive number is the index (0-based) of assigned gt.
        The assignment is done in following steps, the order matters.

        1. assign every bbox to the background
        2. assign proposals whose iou with all gts < neg_iou_thr to 0
        3. for each bbox, if the iou with its nearest gt >= pos_iou_thr,
           assign it to that bbox
        4. for each gt bbox, assign its nearest proposals (may be more than
           one) to itself

        Args:
            bboxes (Tensor): Bounding boxes to be assigned, shape(n, 4).
            gt_bboxes (Tensor): Groundtruth boxes, shape (k, 4).
            gt_bboxes_ignore (Tensor, optional): Ground truth bboxes that are
                labelled as `ignored`, e.g., crowd boxes in COCO.
            gt_labels (Tensor, optional): Label of gt_bboxes, shape (k, ).

        Returns:
            :obj:`AssignResult`: The assign result.

        Example:
            >>> self = MaxIoUAssigner(0.5, 0.5)
            >>> bboxes = torch.Tensor([[0, 0, 10, 10], [10, 10, 20, 20]])
            >>> gt_bboxes = torch.Tensor([[0, 0, 10, 9]])
            >>> assign_result = self.assign(bboxes, gt_bboxes)
            >>> expected_gt_inds = torch.LongTensor([1, 0])
            >>> assert torch.all(assign_result.gt_inds == expected_gt_inds)
        """
        assign_on_cpu = True if (self.gpu_assign_thr >= 0) and (
            gt_bboxes.shape[0] > self.gpu_assign_thr) else False
        # compute overlap and assign gt on CPU when number of GT is large
        if assign_on_cpu:
            device = bboxes.device
            bboxes = bboxes.cpu()
            gt_bboxes = gt_bboxes.cpu()
            if gt_bboxes_ignore is not None:
                gt_bboxes_ignore = gt_bboxes_ignore.cpu()
            if gt_labels is not None:
                gt_labels = gt_labels.cpu()

        overlaps = self.iou_calculator(gt_bboxes, bboxes, mode=self.assign_metric)

        if (self.ignore_iof_thr > 0 and gt_bboxes_ignore is not None
                and gt_bboxes_ignore.numel() > 0 and bboxes.numel() > 0):
            if self.ignore_wrt_candidates:
                ignore_overlaps = self.iou_calculator(
                    bboxes, gt_bboxes_ignore, mode='iof')
                ignore_max_overlaps, _ = ignore_overlaps.max(dim=1)
            else:
                ignore_overlaps = self.iou_calculator(
                    gt_bboxes_ignore, bboxes, mode='iof')
                ignore_max_overlaps, _ = ignore_overlaps.max(dim=0)
            overlaps[:, ignore_max_overlaps > self.ignore_iof_thr] = -1

        #assign_result =self.lvl_assign_wrt_ranking(gt_bboxes, overlaps, gt_labels)

        assign_result =self.assign_wrt_ranking(overlaps, gt_labels)

        ######### POSTERIOR #############
        device = bboxes.device
        bbox_preds = bbox_preds.to(device)
        cls_scores = cls_scores.to(device)
        bbox_preds = torch.transpose(bbox_preds, 0, 1)

        bbox_preds = self.bbox_coder.decode(bboxes, bbox_preds)

        num_gt = gt_bboxes.size(0)
        num_bboxes = bboxes.size(0)

        # loss-aware assigning within deformable candidate box
        can_positive_mask = assign_result.gt_inds > 0
        can_positive_inds = torch.nonzero(can_positive_mask)
        #print(can_positive_inds)

        poscan = assign_result.gt_inds[can_positive_inds].squeeze(-1)
        can_other_mask = assign_result.gt_inds <= 0

        can_pos_scores = cls_scores[:,can_positive_inds].squeeze(-1)#.to(device)

        can_pos_scores = torch.transpose(can_pos_scores, 0, 1)
        can_bbox_pred = bbox_preds[can_positive_inds,:].squeeze()#.to(device)
        #can_bbox_pred = torch.transpose(can_bbox_pred, 0, 1)

        can_pos_iou = self.iou_calculator(gt_bboxes.to(device), can_bbox_pred, mode ='iou')
        can_pos_iou = can_pos_iou[poscan-1,range(poscan.size(0))]
        can_pos_cls, _ = torch.max(can_pos_scores,1)

        #self.count +=1

        #if self.count % 50 == 0:
        #    print(can_pos_cls)

        can_pos_quality = can_pos_iou + can_pos_cls.sigmoid() * self.cls_w
        can_pos_quality = can_pos_quality.unsqueeze(0).repeat(num_gt, 1) # size of gt, pos anchors

        gt_poscan = torch.zeros_like(can_pos_quality) - 100 # size of gt, pos anchors
        gt_poscan[poscan-1,range(poscan.size(0))] = can_pos_quality[poscan-1,range(poscan.size(0))]

        if self.topq >= can_pos_quality.size(1):
            topq = can_pos_quality.size(1)
        else:
            topq = self.topq
        gt_max_quality, gt_argmax_quality = gt_poscan.topk(topq, dim=1, largest=True, sorted=True)  # gt_argmax_quality [num_gt, q]

        assign_result_pre_gt = assign_result.gt_inds

        assigned_gt_inds_init = assign_result_pre_gt * can_other_mask
        assigned_pos_prior = torch.zeros((num_gt, topq, 5),device=device)
        #assigned_pos_prior_t1 = torch.zeros((num_gt, 1, 5), device=device)

        for i in range(num_gt):
            for j in range(topq):
                index = gt_argmax_quality[i,j]
                remap_inds = can_positive_inds[index,0]
                assigned_gt_inds_init[remap_inds] = assign_result_pre_gt[remap_inds]
                assigned_pos_prior[i,j,:] = bboxes[remap_inds,:] # get the location of assigned positive prior for each gt
                #print(assign_result.gt_inds[can_positive_inds[index,0]])
        assign_result.gt_inds = assigned_gt_inds_init

        ########## Fine-grained instance

        if self.inside_circle == 'circle':
            center_distance = self.iou_calculator(gt_bboxes, bboxes, mode = 'center_distance2')
            width_gt = gt_bboxes[...,2]
            height_gt = gt_bboxes[...,3]
            # scale [0, 32]^2 r=2, scale [32, 256]^2 r= 1.5, scale [256, +inf]^2 r=1 for scale normalization
            '''
            scale = width_gt * height_gt
            scale_1 = scale <= 32*32
            scale_2 = (scale> 32*32) & (scale <= 256*256)
            scale_3 = scale > 256*256
            r = [2, 1.5, 1]
            gt_circle2 = ((width_gt/2)**2 + (height_gt/2) **2) *(scale_1*r[0]**2+scale_2*r[1]**2+scale_3*r[2]**2)
            '''
            r=1.5
            gt_circle = ((width_gt/2)**2 + (height_gt/2) **2) * r * r
            inside_flag = center_distance <= gt_circle[...,None]
            length = range(assign_result.gt_inds.size(0))
            inside_mask = inside_flag[(assign_result.gt_inds-1).clamp(min=0), length]
            assign_result.gt_inds *= inside_mask

        elif self.inside_circle == 'gaussian':
            device1 = gt_bboxes.device
            xy_t, sigma_t = self.xy_wh_r_2_xy_sigma(gt_bboxes)
            xy_a = bboxes[...,:2]
            xy_t = xy_t[...,None,:,:2].unsqueeze(-1)
            xy_a = xy_a[...,:,None,:2].unsqueeze(-1)
            inv_sigma_t = torch.stack((sigma_t[..., 1, 1], -sigma_t[..., 0, 1],
                                      -sigma_t[..., 1, 0], sigma_t[..., 0, 0]),
                                      dim=-1).reshape(-1, 2, 2)
            inv_sigma_t = inv_sigma_t / sigma_t.det().unsqueeze(-1).unsqueeze(-1)
            gaussian = torch.exp(-0.5*(xy_a-xy_t).permute(0, 1, 3, 2).matmul(inv_sigma_t).matmul(xy_a-xy_t)).squeeze(-1).squeeze(-1) #/(2*3.1415926*sigma_t.det())
            inside_flag = gaussian >= torch.exp(torch.tensor([-self.gauss_thr])).to(device1)
            length = range(assign_result.gt_inds.size(0))
            inside_mask = inside_flag[length, (assign_result.gt_inds-1).clamp(min=0)]
            assign_result.gt_inds *= inside_mask

        elif self.inside_circle == 'dgmm':
            device1 = gt_bboxes.device
            xy_gt, sigma_t = self.xy_wh_r_2_xy_sigma(gt_bboxes)
            # get the mean of the positive samples
            pos_prior_mean = torch.mean(assigned_pos_prior[...,:2], dim=-2)
            _, sigma_t = self.xy_wh_r_2_xy_sigma(gt_bboxes)
            xy_pt = pos_prior_mean
            xy_a = bboxes[...,:2]
            xy_gt = xy_gt[...,None,:,:2].unsqueeze(-1)
            xy_pt = xy_pt[...,None,:,:2].unsqueeze(-1)
            xy_a = xy_a[...,:,None,:2].unsqueeze(-1)
            inv_sigma_t = torch.stack((sigma_t[..., 1, 1], -sigma_t[..., 0, 1],
                                      -sigma_t[..., 1, 0], sigma_t[..., 0, 0]),
                                      dim=-1).reshape(-1, 2, 2)
            inv_sigma_t = inv_sigma_t / sigma_t.det().unsqueeze(-1).unsqueeze(-1)
            gaussian_gt = torch.exp(-0.5*(xy_a-xy_gt).permute(0, 1, 3, 2).matmul(inv_sigma_t).matmul(xy_a-xy_gt)).squeeze(-1).squeeze(-1) #/(2*3.1415926*sigma_t.det())
            gaussian_pt = torch.exp(-0.5*(xy_a-xy_pt).permute(0, 1, 3, 2).matmul(inv_sigma_t).matmul(xy_a-xy_pt)).squeeze(-1).squeeze(-1)

            gaussian = 0.7*gaussian_gt + 0.3*gaussian_pt

            inside_flag = gaussian >= torch.exp(torch.tensor([-self.gauss_thr])).to(device1)
            length = range(assign_result.gt_inds.size(0))
            inside_mask = inside_flag[length, (assign_result.gt_inds-1).clamp(min=0)]
            assign_result.gt_inds *= inside_mask

        if assign_on_cpu:
            assign_result.gt_inds = assign_result.gt_inds.to(device)
            assign_result.max_overlaps = assign_result.max_overlaps.to(device)
            if assign_result.labels is not None:
                assign_result.labels = assign_result.labels.to(device)
        return assign_result

    def lvl_assign_wrt_ranking(self,  gt_bboxes, overlaps, gt_labels=None):
        # assign samples inside a certain FPN layer
        num_gts, num_bboxes = overlaps.size(0), overlaps.size(1)

        # 1. assign -1 by default
        assigned_gt_inds = overlaps.new_full((num_bboxes,),
                                             -1,
                                             dtype=torch.long)

        if num_gts == 0 or num_bboxes == 0:
            # No ground truth or boxes, return empty assignment
            max_overlaps = overlaps.new_zeros((num_bboxes,))
            if num_gts == 0:
                # No truth, assign everything to background
                assigned_gt_inds[:] = 0
            if gt_labels is None:
                assigned_labels = None
            else:
                assigned_labels = overlaps.new_full((num_bboxes,),
                                                    -1,
                                                    dtype=torch.long)
            return AssignResult(
                num_gts,
                assigned_gt_inds,
                max_overlaps,
                labels=assigned_labels)

        # for each anchor, which gt best overlaps with it
        # for each anchor, the max iou of all gts
        # assign negative samples

        scale_range = torch.tensor([0,64,128,256,512], device=overlaps.device)
        scale_range = scale_range #** 2
        num_priors = torch.tensor([0,0,0,0,0], device=overlaps.device)
        start_lvl = torch.tensor([0,0,0,0,0], device=overlaps.device)
        end_lvl = torch.tensor([0,0,0,0,0], device=overlaps.device)

        feat_size = 1024/8

        # assign labels according to FPN level

        #overlaps_fpn = overlaps
        overlaps_fpn = torch.zeros_like(overlaps)

        #gt_hbb = self.obb2hbb_le135_xywh(gt_bboxes)

        # assign lvls accodting to its
        #gt_hbb_scales = gt_hbb[:,2]*gt_hbb[:,3] # assign by hbb scales
        #gt_hbb_scales = gt_bboxes[:,2]*gt_bboxes[:,3]  assign by obb scales
        gt_hbb_scales, _ = torch.max(gt_bboxes[:,2:4], 1)

        gt_hbb_scales[(gt_hbb_scales>scale_range[0]) & (gt_hbb_scales<=scale_range[1])] = 0
        gt_hbb_scales[(gt_hbb_scales>scale_range[1]) & (gt_hbb_scales<=scale_range[2])] = 1
        gt_hbb_scales[(gt_hbb_scales>scale_range[2]) & (gt_hbb_scales<=scale_range[3])] = 2
        gt_hbb_scales[(gt_hbb_scales>scale_range[3]) & (gt_hbb_scales<=scale_range[4])] = 3
        gt_hbb_scales[(gt_hbb_scales>scale_range[4])] = 4

        for lvl in range(5):
            start_lvl[lvl] = end_lvl[lvl-1] + num_priors[lvl]

            num_priors[lvl] = feat_size ** 2
            feat_size = feat_size/2

            end_lvl[lvl] = end_lvl[lvl-1] + num_priors[lvl]

            #overlaps_lvl = overlaps[:,start_lvl:end_lvl]
        start_lvl = start_lvl.long()
        end_lvl = end_lvl.long()

        # set other levels to zero, set the assigned level to a positive number
        for i in range(num_gts):
            gt_hbb_lvl = gt_hbb_scales[i].long()
            overlaps_fpn[i,start_lvl[gt_hbb_lvl]:end_lvl[gt_hbb_lvl]] = overlaps[i,start_lvl[gt_hbb_lvl]:end_lvl[gt_hbb_lvl]]

        max_overlaps, argmax_overlaps = overlaps_fpn.max(dim=0)
        # for each gt, topk anchors
        # for each gt, the topk of all proposals
        gt_max_overlaps, gt_argmax_overlaps = overlaps_fpn.topk(self.topk, dim=1, largest=True, sorted=True)  # gt_argmax_overlaps [num_gt, k]

        #print(max_overlaps)

        assigned_gt_inds[(max_overlaps >= 0)
                             & (max_overlaps < 1.0)] = 0

        #assign wrt ranking
        for i in range(num_gts):
            for j in range(self.topk):
                max_overlap_inds = overlaps_fpn[i,:] == gt_max_overlaps[i,j]
                assigned_gt_inds[max_overlap_inds] = i + 1

        if gt_labels is not None:
            assigned_labels = assigned_gt_inds.new_full((num_bboxes,), -1)
            pos_inds = torch.nonzero(
                assigned_gt_inds > 0, as_tuple=False).squeeze()
            if pos_inds.numel() > 0:
                assigned_labels[pos_inds] = gt_labels[
                    assigned_gt_inds[pos_inds] - 1]
        else:
            assigned_labels = None

        return AssignResult(
            num_gts, assigned_gt_inds, max_overlaps, labels=assigned_labels)

    def assign_wrt_ranking(self,  overlaps, gt_labels=None):
        num_gts, num_bboxes = overlaps.size(0), overlaps.size(1)

        # 1. assign -1 by default
        assigned_gt_inds = overlaps.new_full((num_bboxes,),
                                             -1,
                                             dtype=torch.long)

        if num_gts == 0 or num_bboxes == 0:
            # No ground truth or boxes, return empty assignment
            max_overlaps = overlaps.new_zeros((num_bboxes,))
            if num_gts == 0:
                # No truth, assign everything to background
                assigned_gt_inds[:] = 0
            if gt_labels is None:
                assigned_labels = None
            else:
                assigned_labels = overlaps.new_full((num_bboxes,),
                                                    -1,
                                                    dtype=torch.long)
            return AssignResult(
                num_gts,
                assigned_gt_inds,
                max_overlaps,
                labels=assigned_labels)

        # for each anchor, which gt best overlaps with it
        # for each anchor, the max iou of all gts
        max_overlaps, argmax_overlaps = overlaps.max(dim=0)
        # for each gt, topk anchors
        # for each gt, the topk of all proposals
        gt_max_overlaps, gt_argmax_overlaps = overlaps.topk(self.topk, dim=1, largest=True, sorted=True)  # gt_argmax_overlaps [num_gt, k]

        assigned_gt_inds[(max_overlaps >= 0)
                             & (max_overlaps < 0.8)] = 0

        #assign wrt ranking
        for i in range(num_gts):
            for j in range(self.topk):
                max_overlap_inds = overlaps[i,:] == gt_max_overlaps[i,j]
                assigned_gt_inds[max_overlap_inds] = i + 1

        if gt_labels is not None:
            assigned_labels = assigned_gt_inds.new_full((num_bboxes,), -1)
            pos_inds = torch.nonzero(
                assigned_gt_inds > 0, as_tuple=False).squeeze()
            if pos_inds.numel() > 0:
                assigned_labels[pos_inds] = gt_labels[
                    assigned_gt_inds[pos_inds] - 1]
        else:
            assigned_labels = None

        return AssignResult(
            num_gts, assigned_gt_inds, max_overlaps, labels=assigned_labels)

    def xy_wh_r_2_xy_sigma(self, xywhr):
        """Convert oriented bounding box to 2-D Gaussian distribution.

        Args:
            xywhr (torch.Tensor): rbboxes with shape (N, 5).

        Returns:
            xy (torch.Tensor): center point of 2-D Gaussian distribution
                with shape (N, 2).
            sigma (torch.Tensor): covariance matrix of 2-D Gaussian distribution
                with shape (N, 2, 2).
        """
        _shape = xywhr.shape
        assert _shape[-1] == 5
        xy = xywhr[..., :2]
        wh = xywhr[..., 2:4].clamp(min=1e-7, max=1e7).reshape(-1, 2)
        r = xywhr[..., 4]
        cos_r = torch.cos(r)
        sin_r = torch.sin(r)
        R = torch.stack((cos_r, -sin_r, sin_r, cos_r), dim=-1).reshape(-1, 2, 2)
        S = 0.5 * torch.diag_embed(wh)

        sigma = R.bmm(S.square()).bmm(R.permute(0, 2,
                                                1)).reshape(_shape[:-1] + (2, 2))

        return xy, sigma

    def obb2hbb_le135_xywh(self, rotatex_boxes):
        """Convert oriented bounding boxes to horizontal bounding boxes.

        Args:
            obbs (torch.Tensor): [x_ctr,y_ctr,w,h,angle]

        Returns:
            hbbs (torch.Tensor): [x_ctr,y_ctr,w,h,-pi/2]
        """
        polys = self.obb2poly_le135(rotatex_boxes)
        xmin, _ = polys[:, ::2].min(1)
        ymin, _ = polys[:, 1::2].min(1)
        xmax, _ = polys[:, ::2].max(1)
        ymax, _ = polys[:, 1::2].max(1)
        bboxes = torch.stack([xmin, ymin, xmax, ymax], dim=1)
        x_ctr = (bboxes[:, 2] + bboxes[:, 0]) / 2.0
        y_ctr = (bboxes[:, 3] + bboxes[:, 1]) / 2.0
        edges1 = torch.abs(bboxes[:, 2] - bboxes[:, 0])
        edges2 = torch.abs(bboxes[:, 3] - bboxes[:, 1])
        #angles = bboxes.new_zeros(bboxes.size(0))
        inds = edges1 < edges2
        rotated_boxes = torch.stack((x_ctr, y_ctr, edges1, edges2), dim=1)
        rotated_boxes[inds, 2] = edges2[inds]
        rotated_boxes[inds, 3] = edges1[inds]
        #rotated_boxes[inds, 4] = np.pi / 2.0
        return rotated_boxes

    def obb2poly_le135(self, rboxes):
        """Convert oriented bounding boxes to polygons.

        Args:
            obbs (torch.Tensor): [x_ctr,y_ctr,w,h,angle]

        Returns:
            polys (torch.Tensor): [x0,y0,x1,y1,x2,y2,x3,y3]
        """
        N = rboxes.shape[0]
        if N == 0:
            return rboxes.new_zeros((rboxes.size(0), 8))
        x_ctr, y_ctr, width, height, angle = rboxes.select(1, 0), rboxes.select(
            1, 1), rboxes.select(1, 2), rboxes.select(1, 3), rboxes.select(1, 4)
        tl_x, tl_y, br_x, br_y = \
            -width * 0.5, -height * 0.5, \
            width * 0.5, height * 0.5
        rects = torch.stack([tl_x, br_x, br_x, tl_x, tl_y, tl_y, br_y, br_y],
                            dim=0).reshape(2, 4, N).permute(2, 0, 1)
        sin, cos = torch.sin(angle), torch.cos(angle)
        M = torch.stack([cos, -sin, sin, cos], dim=0).reshape(2, 2,
                                                            N).permute(2, 0, 1)
        polys = M.matmul(rects).permute(2, 1, 0).reshape(-1, N).transpose(1, 0)
        polys[:, ::2] += x_ctr.unsqueeze(1)
        polys[:, 1::2] += y_ctr.unsqueeze(1)

        return polys.contiguous()


#!/usr/bin/env python3
# -*- coding:utf-8 -*-
import json
import sys
from utils import background_foreground_color, background_color
# Json valider
# https://jsonlint.com/

class Camera:
    "Camera class"

    def __init__(self, json_drone, json_camera):
        """Camera data """
        self.drone_brand = json_drone['brand']
        self.drone_name = json_drone['name']
        self.camera_id = json_camera['id']
        self.camera_name = json_camera['camera']
        self.camera_resolution_x = json_camera['resolution_X']
        self.camera_resolution_y = json_camera['resolution_Y']
        self.camera_focal = json_camera['focal']
        self.camera_x_sensor_size = json_camera['x_sensor_size']
        self.camera_y_sensor_size = json_camera['y_sensor_size']

    def __str__(self):

        camera_string = F"{self.camera_id :^12s}|{self.drone_brand+' '+self.drone_name :^22s}|{self.camera_name :^20s}|"
        camera_string += F"{self.camera_resolution_x:>9} x {self.camera_resolution_y:<8}|"
        camera_string += F"{self.camera_focal:^13}|"
        camera_string += F"{self.camera_x_sensor_size:>5} x {self.camera_y_sensor_size:<6}|"
        return camera_string

    def header(self):

        # Header
        header = f"{'Id':^12s}|{'Drone':^22s}|{'Camera':^20s}|"
        header += f"{'Sensor resolution px':^17s}|"
        header += f"{'Focal mm':^13s}|"
        header += f"{'Sensor size mm':^13s}|"
        return header

class Drones:
    "Drones class"

    def __init__(self, json_file='drones.json'):
        """Drones get data from json file """
        self.camera_dict = {}
        # Read JSON file
        with open(json_file) as data_file:
            self.data_loaded = json.load(data_file)

        for drone in self.data_loaded['Drones']:
            for cam in drone['Cameras']:
                camdrone = Camera(drone, cam)
                self.camera_dict[camdrone.camera_id] = Camera(drone, cam)

    def print(self):
        """Custom print() function."""

        # Header
        header = f"{'Id':^12s}|{'Drone':^22s}|{'Camera':^20s}|"
        header += f"{'Sensor resolution px':^17s}|"
        header += f"{'Focal mm':^13s}|"
        header += f"{'Sensor size mm':^13s}|"
        print(background_foreground_color(header, 93, 11))

        # Cameras with alternate colors
#        for i,camdrone in enumerate(self.camera_dict.items()):
        for i, (key, camdrone) in enumerate(self.camera_dict.items()):
            if i % 2 == 0:
                print(background_color(camdrone.__str__(), 28))
            else:
                print(background_color(camdrone.__str__(), 19))

    def get_camera_by_id(self, camera_id):
        """return camera with camera_number index
        """
        cam = self.camera_dict.get(camera_id)
        if cam:
            return cam
        else:
            print(F"camera_id {background_color(camera_id,9)} is not in the database. Available camera are")
            self.print()
            sys.exit(-1)

def main():
    """la fonction main"""
    drones_db = Drones()
    drones_db.print()

if __name__ == '__main__':
    main()