# Copyright (C) 2020 Arm Limited or its affiliates. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the License); you may
# not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an AS IS BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Description:
# The scheduler costs various strategies for scheduling the network in order to select the block configuration.
import copy
import enum
from functools import lru_cache

import numpy as np

from . import live_range
from . import npu_performance
from . import stats_writer
from .data_type import DataType
from .high_level_command_stream_generator import calc_allowed_ofm_ifm_overlap_for_pass_list
from .nn_graph import CascadedPass
from .nn_graph import PassPlacement
from .nn_graph import SchedulerRewrite
from .nn_graph import SchedulingStrategy
from .npu_performance import make_bandwidth_array
from .npu_performance import make_cycles_array
from .npu_performance import make_macs_array
from .npu_performance import make_metrics_arrays
from .npu_performance import PassCycles
from .numeric_util import full_shape
from .operation import NpuBlockType
from .operation import Op
from .operation import Operation
from .shared_buffer_allocation import find_block_configs_suitable_for_pass_and_shared_buffer
from .shared_buffer_allocation import shared_buffer_allocation_for_pass_and_block_config
from .tensor import MemArea
from .tensor import MemType
from .tensor import TensorFormat
from .tensor import TensorPurpose
from .tensor import TensorSubPurpose

class ParetoMetric(enum.Enum):
    BwCycMem = 1
    BwCycMemBlkH = 2

    def __str__(self):
        return self.name

class SchedulerOptions:
    def __init__(
        self,
        use_cascading=True,
        verbose_schedule=False,
        verbose_pareto_frontier_schedules=False,
        use_ifm_streaming=True,
        pareto_metric=ParetoMetric.BwCycMem,
        use_nhcwb16_between_cascaded_passes=True,
        keep_scale_placement=False,
    ):
        self.use_cascading = use_cascading
        self.verbose_schedule = verbose_schedule
        self.verbose_pareto_frontier_schedules = verbose_pareto_frontier_schedules
        self.use_ifm_streaming = use_ifm_streaming
        self.pareto_metric = pareto_metric
        self.use_nhcwb16_between_cascaded_passes = use_nhcwb16_between_cascaded_passes
        self.keep_scale_placement = keep_scale_placement

    def __str__(self):
        return type(self).__name__ + ": " + str(self.__dict__)

    __repr__ = __str__

class Strategy:
    __slots__ = "strat", "param", "passes", "block_configs", "rewrite_list", "bws", "macs", "cycles", "sram_used"

    def __init__(self, strat, param, passes, block_configs, rewrite_list, bws, macs, cycles, sram_used):
        self.strat = strat
        self.param = param
        self.passes = passes
        self.block_configs = block_configs
        self.rewrite_list = (
            rewrite_list  # list of (SchedulerRewrite, Tensor, new sub purpose, purpose param a, purpose param b, pass)
        )
        self.bws = bws
        self.macs = macs
        self.cycles = cycles
        self.sram_used = sram_used

    def __eq__(self, other):
        if self.strat != other.strat:
            return False
        if self.param != other.param:
            return False
        if self.block_configs != other.block_configs:
            return False
        if self.passes != other.passes:
            return False
        if (self.bws != other.bws).any():
            return False
        if (self.macs != other.macs).any():
            return False
        if (self.cycles != other.cycles).any():
            return False
        if self.sram_used != other.sram_used:
            return False
        return True

    def empty(self):
        return not self.passes

    def key(self):
        return self.passes[-1]

    def clone(self):
        return Strategy(
            self.strat,
            self.param,
            self.passes,
            self.block_configs,
            self.rewrite_list,
            self.bws,
            self.macs,
            self.cycles,
            self.sram_used,
        )

    def __str__(self):
        return "<scheduler.Strategy: %s %s %s %s %s %s %s>" % (
            self.strat,
            self.passes,
            self.rewrite_list,
            self.bws,
            self.macs,
            self.cycles,
            self.sram_used,
        )

    __repr__ = __str__

class StrategySet:
    __slots__ = "strats", "bws", "macs", "cycles", "max_sram_used", "total_sram_used"

    def __init__(self, strats=None):
        if strats is None:
            strats = dict()
        self.strats = strats  # final pass in packed pass -> Strategy
        self.bws, self.macs, self.cycles = make_metrics_arrays()
        self.max_sram_used = 0
        self.total_sram_used = 0

    def update_statistics(self):
        self.bws = make_bandwidth_array()
        self.max_sram_used = 0
        for ps, strat in self.strats.items():
            self.bws += strat.bws
            self.macs += strat.macs
            self.cycles += strat.cycles
            self.max_sram_used = max(self.max_sram_used, strat.sram_used)
            self.total_sram_used += strat.sram_used

    def clone_add_strategy(self, new_strat):
        key = new_strat.key()
        if key in self.strats:
            assert new_strat == self.strats[key]
            return self
        else:
            new_strats = dict(self.strats)
            new_strats[key] = new_strat
            new_set = StrategySet(new_strats)
            new_set.bws = self.bws + new_strat.bws
            new_set.macs = self.macs + new_strat.macs
            new_set.cycles = self.cycles + new_strat.cycles
            new_set.max_sram_used = max(self.max_sram_used, new_strat.sram_used)
            new_set.total_sram_used = self.total_sram_used + new_strat.sram_used
            return new_set

    def __eq__(self, other):
        if (self.bws != other.bws).any():
            return False
        if (self.macs != other.macs).any():
            return False
        if (self.cycles != other.cycles).any():
            return False
        if self.max_sram_used != other.max_sram_used:
            return False
        if self.total_sram_used != other.total_sram_used:
            return False
        if self.strats != other.strats:
            return False
        return True

    def __str__(self):
        return "<scheduler.StrategySet: max_sram_used=%s passes_covered=%s>" % (
            self.max_sram_used,
            list(ps.name for ps in self.strats),
        )

    __repr__ = __str__

empty_strategy = Strategy(
    SchedulingStrategy.Unknown, None, [], [], [], make_bandwidth_array(), make_macs_array(), make_cycles_array(), 0
)
INFINITY = 1e30

ABORT_SEARCH = []

def flatten_list_of_lists(lstlst):
    lst = []
    for v in lstlst:
        lst.extend(v)
    return lst

class DynamicProgrammingScheduler:
    def __init__(self, nng, sg, arch, sram_limit, options: SchedulerOptions):
        self.nng = nng
        self.sg = sg
        self.arch = arch
        self.sram_limit = sram_limit
        self.options = copy.copy(options)
        self.use_cascading = options.use_cascading

        if self.arch.feature_map_storage_mem_area != MemArea.Sram:
            self.use_ifm_ofm_overlap = False  # force off IFM/OFM overlap if IFMs and OFMs are not in the SRAM
        else:
            self.use_ifm_ofm_overlap = True

        self.verbose_schedule = options.verbose_schedule
        self.verbose_pareto_frontier_schedules = options.verbose_pareto_frontier_schedules
        self.mem_area = MemArea.Sram

        self.bandwidth_weights = arch.bandwidth_weights
        self.cycles_weight = arch.cycles_weight
        self.max_sram_used_weight = arch.max_sram_used_weight

        self.n_combinations_searched = 0

        self.pareto_max_candidates = 16

        self.ifm_stream_npu_blocks = set(
            (NpuBlockType.ConvolutionMxN, NpuBlockType.ConvolutionDepthWise, NpuBlockType.Pooling,)
        )

    num_pareto_metrics = 4
    view_values = ",".join(["d"] * num_pareto_metrics)
    order_values = ["f%d" % (idx,) for idx in range(num_pareto_metrics)]

    def pareto_metric(self, candidate):
        strat, strat_set = candidate
        total_cycles = strat.cycles[PassCycles.Total] + strat_set.cycles[PassCycles.Total]
        bws = strat.bws + strat_set.bws
        last_block_height = 0
        if self.options.pareto_metric == ParetoMetric.BwCycMemBlkH and len(strat.block_configs) > 0:
            last_block_height = strat.block_configs[-1][0]

        return (
            np.tensordot(bws, self.bandwidth_weights, axes=3) + total_cycles * self.cycles_weight,
            strat_set.max_sram_used,
            strat.sram_used,
            last_block_height,
        )

    def filter_pareto_frontier(self, candidates, remove_equally_good_candidates):

        candidates = [cand for cand in candidates if max(cand[0].sram_used, cand[1].max_sram_used) <= self.sram_limit]

        if len(candidates) <= 1:
            return candidates
        assert remove_equally_good_candidates
        pareto_vals = np.zeros((len(candidates), DynamicProgrammingScheduler.num_pareto_metrics))
        ids = np.arange(len(candidates), dtype=np.int32)
        for idx, cand in enumerate(candidates):
            pareto_vals[idx] = self.pareto_metric(cand)

        sort_order = np.argsort(
            pareto_vals.view(DynamicProgrammingScheduler.view_values),
            order=DynamicProgrammingScheduler.order_values,
            axis=0,
            kind="stable",
        ).flatten()
        pareto_vals = pareto_vals[sort_order]
        ids = ids[sort_order]

        pareto_frontier = []
        while len(ids) > 0:
            pareto_frontier.append(candidates[ids[0]])
            not_dominated_by_first = (pareto_vals < pareto_vals[0]).any(axis=1)
            ids = ids[not_dominated_by_first]
            pareto_vals = pareto_vals[not_dominated_by_first]

        if len(pareto_frontier) > self.pareto_max_candidates:
            pareto_frontier = self.sort_by_candidate_metric(pareto_frontier)
            pareto_frontier = pareto_frontier[: self.pareto_max_candidates]

        return pareto_frontier

    def candidate_metric(self, candidate):
        strat, strat_set = candidate
        max_sram_used = max(strat_set.max_sram_used, strat.sram_used)
        bws = strat.bws + strat_set.bws
        total_cycles = strat.cycles[PassCycles.Total] + strat_set.cycles[PassCycles.Total]

        return (
            max_sram_used * self.max_sram_used_weight
            + np.tensordot(bws, self.bandwidth_weights, axes=3)
            + total_cycles * self.cycles_weight
        )

    def sort_by_candidate_metric(self, candidate_list):
        sorted_list = list(sorted(candidate_list, key=self.candidate_metric))
        return sorted_list

    def best_candidate(self, candidate_list):
        if len(candidate_list) == 0:
            return ABORT_SEARCH
        if len(candidate_list) == 1:
            return candidate_list[0]
        sorted_list = self.sort_by_candidate_metric(candidate_list)
        return sorted_list[0]

    def graduate_strat(self, strat_type, sram_used, old_strat_data):
        res = []
        for old_strat, old_strat_set in old_strat_data:
            if old_strat.sram_used + sram_used > self.sram_limit:
                continue  # This strategy is bad, drop it
            if old_strat_set.max_sram_used > self.sram_limit:
                continue  # This strategy is bad, drop it
            assert old_strat.strat == SchedulingStrategy.Unknown

            new_strat = old_strat.clone()
            new_strat.strat = strat_type
            new_strat.sram_used = old_strat.sram_used + sram_used

            if self.use_ifm_ofm_overlap:
                overlap = calc_allowed_ofm_ifm_overlap_for_pass_list(
                    new_strat.strat, new_strat.passes, new_strat.block_configs
                )
                new_strat.sram_used -= overlap

            new_strat_set = old_strat_set.clone_add_strategy(new_strat)
            res.append((empty_strategy, new_strat_set))
        return self.filter_pareto_frontier(res, remove_equally_good_candidates=True)

    def append_sram(self, sram_used, old_strat_data):
        res = []
        for old_strat, strat_set in old_strat_data:
            assert old_strat.strat == SchedulingStrategy.Unknown
            assert old_strat.sram_used == 0
            new_strat = old_strat.clone()
            new_strat.sram_used = old_strat.sram_used + sram_used

            res.append((new_strat, strat_set))
        return res

    def append_sram_block_config_performance_metrics(self, sram_used, block_config, metrics, old_strat_data):
        res = []
        for old_strat, strat_set in old_strat_data:
            assert old_strat.strat == SchedulingStrategy.Unknown
            new_strat = old_strat.clone()
            bws, macs, cycles = metrics[:3]

            new_strat.sram_used = old_strat.sram_used + sram_used
            new_strat.block_configs = old_strat.block_configs + [block_config]
            new_strat.bws = old_strat.bws + bws
            new_strat.macs = old_strat.macs + macs
            new_strat.cycles = old_strat.cycles + cycles
            new_strat.bws, new_strat.macs, new_strat.cycles = npu_performance.collate_stats_for_cascaded_pass(
                self.arch, new_strat.bws, new_strat.macs, new_strat.cycles
            )

            res.append((new_strat, strat_set))
        return res

    def append_sram_pass_block_config_performance_metrics_rewrite_list(
        self, sram_used, new_pass, block_config, metrics, rewrite_list, old_strat_data
    ):
        res = []
        for old_strat, strat_set in old_strat_data:
            assert old_strat.strat == SchedulingStrategy.Unknown
            new_strat = old_strat.clone()
            bws, macs, cycles = metrics[:3]
            new_strat.sram_used = old_strat.sram_used + sram_used
            new_strat.block_configs = old_strat.block_configs + [block_config]
            new_strat.bws = old_strat.bws + bws
            new_strat.macs = old_strat.macs + macs
            new_strat.cycles = old_strat.cycles + cycles
            new_strat.passes = old_strat.passes + [new_pass]
            new_strat.bws, new_strat.macs, new_strat.cycles = npu_performance.collate_stats_for_cascaded_pass(
                self.arch, new_strat.bws, new_strat.macs, new_strat.cycles
            )
            new_strat.rewrite_list = old_strat.rewrite_list + rewrite_list
            res.append((new_strat, strat_set))
        return res

    def append_sram_rewrite_list(self, sram_used, rewrite_list, old_strat_data):
        res = []
        for old_strat, strat_set in old_strat_data:
            assert old_strat.strat == SchedulingStrategy.Unknown
            new_strat = old_strat.clone()
            new_strat.sram_used = old_strat.sram_used + sram_used
            new_strat.rewrite_list = old_strat.rewrite_list + rewrite_list
            res.append((new_strat, strat_set))
        return res

    def pass_to_strat(self, strat_data):
        res = {}
        for strat in strat_data[1].strats.values():
            for ps in strat.passes:
                res[ps] = strat
        return res

    def compatible_strats(self, a, b):
        intersection = a.keys() & b.keys()
        for k in intersection:
            if a[k] != b[k]:
                return False
        return True

    def collate_strats_for_passes(self, all_passes):
        if len(all_passes) == 0:
            return [(empty_strategy, StrategySet(dict()))]
        if len(all_passes) == 1:
            return all_passes[0]  # save some space in the common case
        all_strands = [[self.pass_to_strat(strat_data) for strat_data in strand] for strand in all_passes]
        prev_combos = [dict()]
        for j, strand in enumerate(all_strands):
            new_combos = []
            for i, alt in enumerate(strand):
                for prev in prev_combos:
                    if self.compatible_strats(prev, alt):
                        cmb = dict(prev)
                        cmb.update(all_passes[j][i][1].strats)
                        new_combos.append(cmb)
            prev_combos = new_combos

        res = []
        for d in prev_combos:
            s = StrategySet(d)
            s.update_statistics()
            res.append((empty_strategy, s))
        return res

    def search_all_but_one_predecessor(self, ps, pred_pass, pred_pass_data):
        # get the rest of the predecessors
        other_predecessors = [pred for pred in ps.dag_predecessors if pred != pred_pass]
        other_predecessor_data = self.search_pass_list(other_predecessors)

        # pred strat data has an incomplete strategy, which we need
        # to continue on, whereas the other ones have completed strategies.
        # we need to merge these, but keep the incomplete strategy too.

        res = []
        for pred_pass_strat, pred_pass_strat_set in pred_pass_data:
            all_strats = [
                [(empty_strategy, pred_pass_strat_set)],  # pred strat data but with a dummy empty strategy
                other_predecessor_data,  # this one is fine to use as-is
            ]
            collated_strat_data = self.collate_strats_for_passes(all_strats)
            strat_data = [(pred_pass_strat, strat_set) for _, strat_set in collated_strat_data]
            res.extend(strat_data)
        return res

    def calc_non_local_mem_usage(self):
        ignore_subgraph_input_output_tensors = self.sg.placement == PassPlacement.Cpu
        range_set = live_range.extract_live_ranges_from_passes(
            self.sg, self.mem_area, ignore_subgraph_input_output_tensors=ignore_subgraph_input_output_tensors,
        )
        range_dict = range_set.ranges

        # find which ranges overlap passes but aren't input/outputs of the passes.
        # these won't be counted by the dynamic programming search and must be counted in manually.
        end_pos = max(ps.time for ps in self.sg.passes) + 2
        mem_usage = np.zeros(end_pos) + self.sg.base_sram_used
        non_local_mem_usage = np.zeros(end_pos, dtype=np.int64)

        for tens, rng in range_dict.items():
            storage_size = tens.storage_size()
            assert tens.mem_area == self.mem_area
            mem_usage[rng.start_time : rng.end_time] += storage_size

        for ps in self.sg.passes:
            local_mem_usage = 0
            for tens in ps.inputs + ps.outputs + ps.intermediates:
                if tens.mem_area != self.mem_area:
                    continue

                local_mem_usage += tens.storage_size()

            non_local_mem_usage[ps.time] = mem_usage[ps.time] - local_mem_usage

        self.non_local_mem_usage = non_local_mem_usage

    def search(self):
        self.calc_non_local_mem_usage()
        starting_passes = [ps for ps in self.sg.passes if not ps.successors]
        strat_data = self.search_pass_list(starting_passes)

        _, best_set = self.best_candidate(strat_data)

        if self.verbose_pareto_frontier_schedules:
            print(
                "Scheduler searched %d combinations and found %d candidate schedules along the pareto frontier"
                % (self.n_combinations_searched, len(strat_data))
            )
            for idx, (_, strat_set) in enumerate(strat_data):
                extra = ""
                if strat_set == best_set:
                    extra = "(Best candidate)"
                print("Candidate", idx, extra)
                memory_used = {MemArea.Sram: strat_set.max_sram_used}
                stats_writer.print_performance_metrics_for_strat(
                    self.arch,
                    "",
                    strat_set.cycles,
                    strat_set.macs,
                    strat_set.bws,
                    self.nng.batch_size,
                    memory_used,
                    len(self.sg.passes),
                    len(strat_set.strats),
                )

        return best_set

    def search_pass_list(self, pass_list):
        all_strats = []
        for ps in pass_list:
            strat = self.search_output(ps)
            all_strats.append(strat)
        strat_data = self.collate_strats_for_passes(all_strats)
        for strd in strat_data:
            for ps in pass_list:
                assert ps in strd[1].strats  # should have strategies for everything we asked to search
        return strat_data

    def search_predecessors(self, ps):

        # protect against graphs with loops. collate_strats_for_passes will sort this out later so that
        # we have strats for all passes

        pass_list = ps.dag_predecessors
        strat_data = self.search_pass_list(pass_list)

        return strat_data

    @lru_cache(maxsize=None)
    def search_output(self, ps):

        assert ps in self.sg.passes
        candidate_list = []

        candidate_list.extend(self.search_weight_streaming_output(ps))

        if self.options.use_ifm_streaming:
            candidate_list.extend(self.search_ifm_streaming_output(ps))

        best = self.filter_pareto_frontier(candidate_list, remove_equally_good_candidates=True)

        if not best:
            print(
                "Warning: Dynamic search programming algorithm failed for pass %s, invoking fallback strategy"
                % (ps.name,)
            )
            return self.search_predecessors(ps)

        return best

    def search_ifm_streaming_output(self, ps):
        if ps.placement != PassPlacement.Npu:
            return ABORT_SEARCH
        if ps.npu_block_type not in self.ifm_stream_npu_blocks:
            return ABORT_SEARCH
        strat_data = self.search_ifm_streaming_body(ps, False)

        sram_used = self.non_local_mem_usage[ps.time]
        for tens in ps.outputs:
            if tens.mem_area == self.mem_area:
                sram_used += tens.storage_size()

        return self.graduate_strat(SchedulingStrategy.IfmStream, sram_used, strat_data)

    @lru_cache(maxsize=None)
    def search_ifm_streaming_body(self, ps, force_outputs_to_fast_storage):
        if ps.placement != PassPlacement.Npu:
            return ABORT_SEARCH
        if ps.npu_block_type not in self.ifm_stream_npu_blocks:
            return ABORT_SEARCH
        ifm_input_search_resuls = self.search_ifm_streaming_input(ps)
        res = []

        base_sram_used = 0
        for tens in ps.intermediates:
            if tens.mem_area == self.mem_area:
                if tens.purpose == TensorPurpose.Weights:
                    base_sram_used = tens.storage_size(self.arch.weight_estimation_scaling)
                else:
                    base_sram_used += tens.storage_size()

        all_block_configs = self.get_block_configs(ps)
        for block_config in all_block_configs:
            all_strats = []

            if self.use_cascading:
                all_strats.extend(self.search_ifm_streaming_partial(ps, block_config))

            all_strats.extend(ifm_input_search_resuls)

            rewrite_list = []
            sram_used = base_sram_used

            metrics = npu_performance.performance_metrics_for_pass(
                self.arch,
                ps,
                block_config,
                rewrite_list=rewrite_list,
                force_outputs_to_fast_storage=force_outputs_to_fast_storage,
            )

            res.extend(
                self.append_sram_pass_block_config_performance_metrics_rewrite_list(
                    sram_used, ps, block_config, metrics, rewrite_list, all_strats
                )
            )

        self.n_combinations_searched += len(res)
        res = self.filter_pareto_frontier(res, remove_equally_good_candidates=True)
        return res

    def avoid_for_cascading(self, pred_candidate):
        for op in pred_candidate.ops:
            if (
                op.type == Op.ConcatSliceWrite
                and self.arch.feature_map_storage_mem_area != self.arch.fast_storage_mem_area
            ):
                # For SRAM spilling, concat op is avoided as predecessor
                return True
            if len(op.outputs) > 1 or len(op.outputs[0].consumer_list) > 1:
                # The op has consumers in other subgraphs
                return True
        return False

    def search_ifm_streaming_partial(self, ps, block_config):
        if ps.placement != PassPlacement.Npu:
            return ABORT_SEARCH

        if len(ps.inputs) < 1:
            return ABORT_SEARCH

        ifm_tensor = ps.ifm_tensor

        if ifm_tensor is None:
            return ABORT_SEARCH
        if ifm_tensor.purpose != TensorPurpose.FeatureMap:
            return ABORT_SEARCH
        if not ifm_tensor.storage_shape or len(ifm_tensor.storage_shape) != 4:
            return ABORT_SEARCH

        pred_pass_list = []
        for pred_candidate in ps.dag_predecessors:
            if len(pred_candidate.outputs) == 1 and pred_candidate.outputs[0] == ifm_tensor:
                # we found a predecessor that produces this IFM tensor
                if not ifm_tensor.avoid_NHCWB16:
                    # and NHCWB16 format is not to be avoided
                    if len(pred_candidate.successors) == 1 and pred_candidate.successors[0] == ps:
                        # and it only has one successor, namely us
                        if pred_candidate.placement == PassPlacement.Npu:
                            if pred_candidate.npu_block_type in self.ifm_stream_npu_blocks:
                                # and it is on the Npu
                                if not self.avoid_for_cascading(pred_candidate):
                                    # and fusable - it's a candidate
                                    pred_pass_list.append(pred_candidate)

        if not pred_pass_list:
            return ABORT_SEARCH

        all_candidates = []
        for pred_pass in pred_pass_list:
            # recurse into the next pass
            ifm_strat_data = self.search_ifm_streaming_body(pred_pass, self.arch.is_spilling_enabled())

            strat_data = self.search_all_but_one_predecessor(ps, pred_pass, ifm_strat_data)
            for strat_opt in strat_data:

                pred_pass_block_config = strat_opt[0].block_configs[-1]
                rolling_buffer_dims = npu_performance.rolling_buffer_dims_from_passes(
                    self.arch, pred_pass, pred_pass_block_config, ps, block_config
                )
                if rolling_buffer_dims is None:
                    continue  # this does not pack properly, skip it.

                sram_used = 0
                for tens in ps.inputs:
                    if tens != ifm_tensor:
                        if tens.mem_area == self.mem_area:
                            sram_used += tens.storage_size()

                rolling_buffer_y, rolling_buffer_x = rolling_buffer_dims

                rewrite_list = [
                    (
                        SchedulerRewrite.ChangeTensorSubPurpose,
                        ifm_tensor,
                        TensorSubPurpose.RollingBufferY,
                        rolling_buffer_y,
                        None,
                        ps,
                    )
                ]
                sram_used += ifm_tensor.storage_size_for_sub_purpose(
                    self.arch, TensorSubPurpose.RollingBufferY, rolling_buffer_y, None
                )

                all_candidates.extend(self.append_sram_rewrite_list(sram_used, rewrite_list, [strat_opt]))

        self.n_combinations_searched += len(all_candidates)
        return all_candidates

    def get_block_configs(self, ps):
        if ps.placement != PassPlacement.Npu:
            return [(1, 1, 1, 1)]  # default

        block_configs = find_block_configs_suitable_for_pass_and_shared_buffer(self.arch, ps)

        # Take a limited number of the largest blocks
        if self.arch.block_config_limit > 0:
            # Sort by block area, followed by depth
            block_configs.sort(key=lambda cfg: (cfg[0] * cfg[1]) << 8 | cfg[3], reverse=True)
            bound = min(len(block_configs), self.arch.block_config_limit)
            # We take 'n' from the fat end of the list, and 'n' from the thin end of the list.
            tmp = block_configs[:bound]
            tmp.extend(block_configs[max(bound, len(block_configs) - bound) :])
            block_configs = tmp

        return block_configs

    def search_ifm_streaming_input(self, ps):
        sram_used = 0
        for tens in ps.inputs:
            if tens.mem_area == self.mem_area:
                sram_used += tens.storage_size()

        return self.append_sram(sram_used, self.search_predecessors(ps))

    def search_weight_streaming_output(self, ps):
        strat_data = self.search_weight_streaming_body(ps)

        sram_used = self.non_local_mem_usage[ps.time]
        for tens in ps.outputs:
            if tens.mem_area == self.mem_area:
                sram_used += tens.storage_size()

        return self.graduate_strat(SchedulingStrategy.WeightStream, sram_used, strat_data)

    @lru_cache(maxsize=None)
    def search_weight_streaming_body(self, ps):

        strat_data = self.search_weight_streaming_input(ps)

        res = []

        all_block_configs = self.get_block_configs(ps)

        for block_config in all_block_configs:

            sram_used = 0
            rewrite_list = []

            for tens in ps.intermediates:
                if tens.mem_area == self.mem_area:
                    if tens.purpose == TensorPurpose.Weights:
                        sram_used += tens.storage_size_for_sub_purpose(
                            self.arch, TensorSubPurpose.DoubleBuffer, block_config[3]
                        )
                        rewrite_list.append(
                            (
                                SchedulerRewrite.ChangeTensorSubPurpose,
                                tens,
                                TensorSubPurpose.DoubleBuffer,
                                block_config[3],
                                None,
                                ps,
                            )
                        )
                    else:
                        sram_used += tens.storage_size()

            metrics = npu_performance.performance_metrics_for_pass(
                self.arch, ps, block_config, rewrite_list=rewrite_list
            )

            res.extend(
                self.append_sram_pass_block_config_performance_metrics_rewrite_list(
                    sram_used, ps, block_config, metrics, rewrite_list, strat_data
                )
            )

        self.n_combinations_searched += len(res)
        res = self.filter_pareto_frontier(res, remove_equally_good_candidates=True)
        return res

    def search_weight_streaming_input(self, ps):
        sram_used = 0
        for tens in ps.inputs:
            if tens.mem_area == self.mem_area:
                sram_used += tens.storage_size()

        return self.append_sram(sram_used, self.search_predecessors(ps))

    def apply_result(self, strat_set, arch):
        pass_to_cascaded_pass = dict()
        for _, strat in strat_set.strats.items():
            # rewrite the tensors that need this first. e.g. make rolling buffers
            inputs = []
            intermediates = []
            outputs = []

            for ps in strat.passes:
                inputs += ps.inputs
                intermediates += ps.intermediates
                outputs += ps.outputs

            for tens in set(inputs) & set(outputs):
                # tensors that are in both sets are intermediates

                # find pass with input/output tensor, and check if they are both placed on NPU
                input_placement = None
                output_placement = None
                for ps in strat.passes:
                    if tens in ps.inputs:
                        input_placement = ps.placement
                    if tens in ps.outputs:
                        output_placement = ps.placement
                if input_placement == output_placement == PassPlacement.Npu:
                    tens.set_format(TensorFormat.NHCWB16, arch)

                intermediates.append(tens)
                inputs.remove(tens)
                outputs.remove(tens)

            for rewrite_op, tens, sub_purpose, param_a, param_b, ps in strat.rewrite_list:
                if rewrite_op == SchedulerRewrite.ChangeTensorSubPurpose:
                    tens.mem_area = self.arch.fast_storage_mem_area
                    tens.mem_type = MemType.Scratch_fast
                    tens.set_new_sub_purpose(sub_purpose, param_a, param_b)
                else:
                    assert 0, "unknown rewrite_op " + str(rewrite_op)

            is_element_wise = True
            for ps in strat.passes:
                assert ps.placement == strat.passes[0].placement
                if not ps.is_element_wise:
                    is_element_wise = False
                    break

            cascaded_pass = CascadedPass(
                strat.passes[0].name,
                strat.strat,
                inputs,
                intermediates,
                outputs,
                strat.passes,
                strat.passes[0].placement,
                is_element_wise,
            )
            assert strat.sram_used >= 0
            cascaded_pass.sram_used = strat.sram_used

            for idx, ps in enumerate(strat.passes):
                assert ps not in pass_to_cascaded_pass
                pass_to_cascaded_pass[ps] = cascaded_pass
                ps.cascade = cascaded_pass
                ps.block_config = strat.block_configs[idx]

                if ps.placement == PassPlacement.Npu:
                    ps.shared_buffer = shared_buffer_allocation_for_pass_and_block_config(
                        self.arch, ps, ps.block_config
                    )
                    assert ps.shared_buffer is not None

                sram_used = max(self.non_local_mem_usage[ps.time], 0)
                for op in ps.ops:
                    subgraph = op.attrs.get("subgraph")
                    if subgraph:
                        subgraph.base_sram_used = sram_used

        # all passes should have a cascaded pass now
        if len(pass_to_cascaded_pass) != len(self.sg.passes):
            print(
                "mismatch: we have %d passes, but only %d have cascaded passes associated"
                % (len(self.sg.passes), len(pass_to_cascaded_pass))
            )
            for ps in self.sg.passes:
                if ps not in pass_to_cascaded_pass:
                    print("%3d pass missing cascaded pass %s" % (ps.time, ps))

            assert len(pass_to_cascaded_pass) == len(self.sg.passes)

        cascaded_passes = []
        if self.sg.placement == PassPlacement.Cpu:
            # Retain the pass order for CPU subgraph
            cascaded_passes = [ps.cascade for ps in self.sg.passes]
        else:
            # we have all the passes, but we need to put them in order and build predecessor/successor links.
            visit_pass_set = set()

            def visit_pass(ps):
                if ps in visit_pass_set:
                    return
                visit_pass_set.add(ps)

                cps = ps.cascade
                dont_traverse = set(cps.passes)

                for ps in cps.passes:
                    for pred in ps.predecessors:
                        if pred in dont_traverse:
                            continue
                        visit_pass(pred)

                cascaded_passes.append(cps)

            starting_passes = [ps for ps in self.sg.passes if not ps.successors]
            for ps in starting_passes:
                visit_pass(ps)

        # reorder so startup init cascaded passes come first
        def is_startup_cascaded_pass(cps):
            if not cps.passes:
                return False
            return cps.placement == PassPlacement.StartupInit

        cascaded_passes = [cps for cps in cascaded_passes if is_startup_cascaded_pass(cps)] + [
            cps for cps in cascaded_passes if not is_startup_cascaded_pass(cps)
        ]

        self.sg.cascaded_passes = cascaded_passes
        self.sg.build_cascaded_pass_links()

        # Check if NHCWB16 and/or fast storage can be used in between cascaded passes
        # (NHCWB16 within cascaded passes has been handled earlier in this function)
        if self.sg.placement == PassPlacement.Npu:
            # Dictionary tensor -> list of ops, containing feature maps that can be attempted
            # to be moved to fast storage
            fast_storage_tensor_rewrites = {}
            last_op_in_subgraph = self.sg.cascaded_passes[-1].passes[-1].primary_op
            # Memory only passes have no primary_op, so use the last op in ops
            if last_op_in_subgraph is None:
                last_op_in_subgraph = self.sg.cascaded_passes[-1].passes[-1].ops[-1]
            for ps in self.sg.cascaded_passes:
                if ps.placement != PassPlacement.Npu:
                    continue
                for output in ps.outputs:
                    if output.purpose != TensorPurpose.FeatureMap or output.avoid_NHCWB16:
                        continue

                    use_NHCWB16 = True
                    use_fast_storage = True
                    rewrites = []
                    for op in output.consumer_list:
                        if op is None:
                            use_NHCWB16 = False
                            use_fast_storage = False
                            continue
                        if op.type == Op.ReduceSum and output.dtype == DataType.int32:
                            use_NHCWB16 = False
                        elif op.type == Op.Reshape:
                            # Using NHCWB16 format for a no-op reshape is only an option if subsequent
                            # consumers do not also need to perform a reshape or if the OFM is going to
                            # be processed by CPU operations. No-op reshape consumers with empty lists
                            # (those that have no consumers, or null-consumers used as list terminators)
                            # must use normal NHWC output.
                            def incompatible_consumers(oper):
                                if oper and oper.type == Op.Reshape:
                                    for consumer in oper.outputs[0].consumer_list:
                                        yield from incompatible_consumers(consumer)
                                yield not oper or not oper.run_on_npu or oper is last_op_in_subgraph

                            if not any(incompatible_consumers(op)):

                                def get_rewrites(oper):
                                    if oper and oper.type == Op.Reshape:
                                        for consumer in oper.outputs[0].consumer_list:
                                            yield from get_rewrites(consumer)
                                        yield oper

                                rewrites.extend(get_rewrites(op))
                                # Detect no-op reshapes by comparing their full input and output tensor shapes.
                                inshape = full_shape(4, op.inputs[0].shape, 1)
                                compatible_shape = [
                                    (inshape == full_shape(4, oper.outputs[0].shape, 1)) for oper in get_rewrites(op)
                                ]
                                use_NHCWB16 = compatible_shape and all(compatible_shape)
                            else:
                                use_NHCWB16 = False
                                use_fast_storage = False
                        use_NHCWB16 &= op.run_on_npu
                        use_fast_storage &= op.run_on_npu

                    if use_fast_storage:
                        fast_storage_tensor_rewrites[output] = rewrites
                    if use_NHCWB16 and self.options.use_nhcwb16_between_cascaded_passes:
                        output.set_format(TensorFormat.NHCWB16, arch)
                        for rewrite_op in rewrites:
                            rewrite_op.outputs[0].set_format(TensorFormat.NHCWB16, arch)
            if arch.is_spilling_enabled():
                # Remember feature maps that can be moved to fast storage for later use
                # in use_fast_storage_for_feature_maps
                self.sg.scheduling_info["feature_map_rewrites"] = fast_storage_tensor_rewrites

def move_scales_to_fast_storage(nng, arch):
    for sg in nng.subgraphs:
        # IFM streamed ops reads bias tensors several times, move these to fast storage
        for cp in sg.cascaded_passes:
            if cp.strategy == SchedulingStrategy.IfmStream:
                # Calculate SRAM usage
                new_size = 0
                all_tens = []
                for ps in cp.passes:
                    pass_tens = np.array([ps.ifm_tensor, ps.ifm2_tensor, ps.ofm_tensor, ps.weight_tensor])
                    pass_tens = np.append(pass_tens, ps.intermediates)
                    for tens in pass_tens:
                        if tens and tens.mem_area == MemArea.Sram and tens not in all_tens:
                            all_tens.append(tens)
                            new_size += tens.storage_size()

                cp.sram_used = new_size

                for ps in cp.passes:
                    if ps.scale_tensor:
                        tens = ps.scale_tensor

                        # Find op using scale tensor
                        op = next((op for op in ps.ops if tens in op.inputs), None)
                        assert op

                        # Create fast storage tensor
                        new_tens = tens.clone_into_fast_storage(arch)
                        new_tens.consumer_list = tens.consumer_list.copy()
                        new_tens.purpose = TensorPurpose.FSBias
                        new_tens_size = new_tens.storage_size()

                        if (cp.sram_used + new_tens_size) <= arch.sram_size:
                            # Create DMA cmd
                            dma_cmd = Operation(Op.DMA, tens.ops[0].name + "_dma")
                            dma_cmd.inputs = [tens]
                            dma_cmd.set_output_tensor(new_tens)
                            dma_cmd.attrs["source"] = tens.mem_area
                            dma_cmd.attrs["destination"] = new_tens.mem_area
                            dma_cmd.run_on_npu = True

                            tens.consumer_list.clear()
                            tens.consumer_list.append(dma_cmd)

                            # Replace tensor and op
                            idx = op.inputs.index(tens)
                            op.inputs[idx] = new_tens

                            ps.ops.insert(0, dma_cmd)
                            ps.scale_tensor = new_tens
                            ps.intermediates.append(new_tens)
                            ps.cascade.intermediates.append(new_tens)

                            cp.sram_used += new_tens_size

def schedule_passes(nng, arch, options: SchedulerOptions):

    for sg in nng.subgraphs:
        sg.base_sram_used = 0

    for sg in nng.subgraphs:
        # re-entering the same nodes from different contexts requires us to
        # build a simplified directed acyclic (DAG) version of the graph to
        # use for traversal, rather than using a visit dictionary. this avoids
        # recursing infinitely due to loops.
        sg.build_pass_dag_predecessors()

        dps = DynamicProgrammingScheduler(nng, sg, arch, arch.sram_size, options)

        strat_set = dps.search()

        dps.apply_result(strat_set, arch)

        if options.verbose_schedule:
            sg.print_cascaded_passes()

def _calc_tens_to_cps(sg, tensor_rewrites):
    # Determines for each tensor the list of affected cascaded passes, in terms of SRAM consumption.
    # Returns dictionary tensor -> list of cascaded passes
    # Note: if cascaded passes are A, B, C, D, and a tensor is output
    # of A and input to D, then it also consumes SRAM in passes B and C.
    if "tens_to_cps" in sg.scheduling_info:
        return sg.scheduling_info["tens_to_cps"]
    # Determine life-time of tensors
    min_index = {}
    max_index = {}
    index = 0
    cps_list = [cps for cps in sg.cascaded_passes if cps.placement == PassPlacement.Npu]
    for cps in cps_list:
        for tens in cps.inputs + cps.outputs:
            if tens in tensor_rewrites:
                min_index[tens] = min(index, min_index.get(tens, len(cps_list)))
                max_index[tens] = index
        index += 1
    # Convert to affected cps-es
    tens_to_cps = {}
    for tens in min_index:
        tens_to_cps[tens] = cps_list[min_index[tens] : max_index[tens] + 1]
    sg.scheduling_info["tens_to_cps"] = tens_to_cps
    return tens_to_cps

def use_fast_storage_for_feature_maps(sg, sram_limit, arch):
    # Attempts to use as much fast storage as possible for feature maps shared between cascaded passes.
    tensor_rewrites = sg.scheduling_info.get("feature_map_rewrites", {})
    tens_to_cps = _calc_tens_to_cps(sg, tensor_rewrites)
    # Sort tensors first on life-time (smallest first), then on size (biggest first)
    tens_list = sorted([(len(tens_to_cps[tens]), -tens.storage_size(), tens.name, tens) for tens in tens_to_cps])
    for _, _, _, tens in tens_list:
        cps_list = tens_to_cps[tens]
        if len(cps_list) < 1:
            continue
        sz = tens.storage_size()
        fits_in_fast_storage = all([cps.sram_used + sz <= sram_limit for cps in cps_list])
        if fits_in_fast_storage:
            tens.mem_area = arch.fast_storage_mem_area
            tens.mem_type = MemType.Scratch_fast
            tens.set_new_sub_purpose(TensorSubPurpose.Standard, None, None)
            assert tens in tensor_rewrites
            # Also rewrite reshapes
            for rewrite_op in tensor_rewrites[tens]:
                tens2 = rewrite_op.outputs[0]
                tens2.mem_area = arch.fast_storage_mem_area
                tens2.mem_type = MemType.Scratch_fast
                tens2.set_new_sub_purpose(TensorSubPurpose.Standard, None, None)
            for cps in cps_list:
                cps.sram_used += sz

def undo_use_fast_storage(sg, arch):
    # Undoes the effects of a previous call to use_fast_storage_for_feature_maps
    tensor_rewrites = sg.scheduling_info.get("feature_map_rewrites", {})
    tens_to_cps = _calc_tens_to_cps(sg, tensor_rewrites)
    mem_area = arch.tensor_storage_mem_area[TensorPurpose.FeatureMap]
    for tens, cps_list in tens_to_cps.items():
        if tens.mem_type == MemType.Scratch_fast:
            sz = tens.storage_size()
            tens.mem_area = mem_area
            tens.mem_type = MemType.Scratch
            # Also undo reshapes
            for rewrite_op in tensor_rewrites[tens]:
                tens2 = rewrite_op.outputs[0]
                tens2.mem_area = mem_area
                tens2.mem_type = MemType.Scratch
            for cps in cps_list:
                cps.sram_used -= sz


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

from ...utils import box_coder_utils, box_utils
from .point_head_template import PointHeadTemplate
from ..backbones_3d.pfe.voxel_set_abstraction import bilinear_interpolate_torch

# from mayavi import mlab

class PointIntraPartOffsetBasedHead(PointHeadTemplate):
    """
    Point-based head for predicting the intra-object part locations.
    Reference Paper: https://arxiv.org/abs/1907.03670
    From Points to Parts: 3D Object Detection from Point Cloud with Part-aware and Part-aggregation Network
    """
    def __init__(self, voxel_size, point_cloud_range, num_class, input_channels, model_cfg, predict_boxes_when_training=False, **kwargs):
        super().__init__(model_cfg=model_cfg, num_class=num_class)
        self.voxel_size = voxel_size
        self.point_cloud_range = point_cloud_range
        self.predict_boxes_when_training = predict_boxes_when_training
        self.cls_layers = self.make_fc_layers(
            fc_cfg=self.model_cfg.CLS_FC,
            input_channels=256,
            output_channels=num_class
        )
        self.part_reg_layers = self.make_fc_layers(
            fc_cfg=self.model_cfg.PART_FC,
            input_channels=256,
            output_channels=3
        )
        target_cfg = self.model_cfg.TARGET_CONFIG
        if target_cfg.get('BOX_CODER', None) is not None:
            self.box_coder = getattr(box_coder_utils, target_cfg.BOX_CODER)(
                **target_cfg.BOX_CODER_CONFIG
            )
            self.box_layers = self.make_fc_layers(
                fc_cfg=self.model_cfg.REG_FC,
                input_channels=256,
                output_channels=self.box_coder.code_size
            )
        else:
            self.box_layers = None

        self.point_feature_layers = self.make_fc_layers(
            fc_cfg=self.model_cfg.FEATURES_FC,
            input_channels=64,
            output_channels=128,
            end_bn_active=True,
        )
        self.pos_encode_layers = self.make_fc_layers(
            fc_cfg=self.model_cfg.POS_FC,
            input_channels=3,
            output_channels=256,
            end_bn_active=True,
        )
        self.encoded_feature_layers = self.make_fc_layers(
            fc_cfg=self.model_cfg.EF_FC,
            input_channels=256,
            output_channels=128,
            end_bn_active=True,
        )
        self.ef_attention_layers = self.make_fc_layers(
            fc_cfg=self.model_cfg.EFA_FC,
            input_channels=256,
            output_channels=1,
        )
        self.pf_attention_layers = self.make_fc_layers(
            fc_cfg=self.model_cfg.PFA_FC,
            input_channels=256,
            output_channels=1,
        )

        self.init_weights()

    def init_weights(self):
        pi = 0.01
        nn.init.constant_(self.cls_layers[6].bias, -np.log((1 - pi) / pi))
        nn.init.normal_(self.box_layers[6].weight, mean=0, std=0.001)
        nn.init.constant_(self.box_layers[6].bias, 0)

    def assign_targets(self, input_dict):
        """
        Args:
            input_dict:
                point_features: (N1 + N2 + N3 + ..., C)
                batch_size:
                point_coords: (N1 + N2 + N3 + ..., 4) [bs_idx, x, y, z]
                gt_boxes (optional): (B, M, 8)
        Returns:
            point_cls_labels: (N1 + N2 + N3 + ...), long type, 0:background, -1:ignored
            point_part_labels: (N1 + N2 + N3 + ..., 3)
        """
        point_coords = input_dict['point_coords']
        # end_point_coords = input_dict['end_point_coords']
        gt_boxes = input_dict['gt_boxes']
        assert gt_boxes.shape.__len__() == 3, 'gt_boxes.shape=%s' % str(gt_boxes.shape)
        assert point_coords.shape.__len__() in [2], 'points.shape=%s' % str(point_coords.shape)

        batch_size = gt_boxes.shape[0]
        extend_gt_boxes = box_utils.enlarge_box3d(
            gt_boxes.view(-1, gt_boxes.shape[-1]), extra_width=self.model_cfg.TARGET_CONFIG.GT_EXTRA_WIDTH
        ).view(batch_size, -1, gt_boxes.shape[-1])
        targets_dict = self.assign_stack_targets(
            points=point_coords,
            gt_boxes=gt_boxes, extend_gt_boxes=extend_gt_boxes,
            set_ignore_flag=True, use_ball_constraint=False,
            ret_part_labels=True, ret_box_labels=(self.box_layers is not None)
        )

        return targets_dict

    def get_loss(self, tb_dict=None):
        tb_dict = {} if tb_dict is None else tb_dict
        point_loss_cls, tb_dict = self.get_cls_layer_loss(tb_dict)
        point_loss_part, tb_dict = self.get_part_layer_loss(tb_dict)
        point_loss = point_loss_cls + point_loss_part

        if self.box_layers is not None:
            point_loss_box, tb_dict = self.get_box_layer_loss(tb_dict)
            point_loss += point_loss_box
        return point_loss, tb_dict

    def interpolate_from_bev_features(self, keypoints, bev_features, batch_size, bev_stride):
        x_idxs = (keypoints[:, 1] - self.point_cloud_range[0]) / self.voxel_size[0]
        y_idxs = (keypoints[:, 2] - self.point_cloud_range[1]) / self.voxel_size[1]
        x_idxs = x_idxs / bev_stride
        y_idxs = y_idxs / bev_stride
        batch_idxs = keypoints[:, 0]

        point_bev_features_list = []
        for k in range(batch_size):
            batch_mask = batch_idxs == k
            cur_x_idxs = x_idxs[batch_mask]
            cur_y_idxs = y_idxs[batch_mask]
            cur_bev_features = bev_features[k].permute(1, 2, 0)  # (H, W, C)
            point_bev_features = bilinear_interpolate_torch(cur_bev_features, cur_x_idxs, cur_y_idxs)
            point_bev_features_list.append(point_bev_features)

        point_bev_features = torch.cat(point_bev_features_list, dim=0)  # (B, N, C0)
        return point_bev_features

    # def show_point_score(self, batch_size, point_coords, cls_scores):
    #     # point_coords = common_utils.get_voxel_centers(
    #     #     sp_tensor.indices[:, 1:], downsample_times=1, voxel_size=self.voxel_size,
    #     #     point_cloud_range=self.point_cloud_range
    #     # )
    #     # point_coords = torch.cat((sp_tensor.indices[:, 0:1].float(), point_coords), dim=1)
    #     # cls_scores = cls_scores[mask]
    #     # point_part_scores = point_part_scores[mask]
    #     for i in range(0, batch_size):
    #         mlab.figure(bgcolor=(0, 0, 0))
    #         mask = point_coords[:, 0] == i
    #         cur_batch_point_coords = point_coords[mask]
    #         cur_x = cur_batch_point_coords[:, 1].cpu().numpy()
    #         cur_y = cur_batch_point_coords[:, 2].cpu().numpy()
    #         cur_z = cur_batch_point_coords[:, 3].cpu().numpy()
    #         # val_depict = (torch.sigmoid(val[mask])).detach().cpu().numpy()
    #         val_depict = (cls_scores[mask]).detach().cpu().numpy()
    #
    #         mlab.points3d(cur_x, cur_y, cur_z, val_depict, scale_factor=0.25, scale_mode='none', mode='sphere',
    #                       line_width=1, colormap='jet')
    #         mlab.colorbar()
    #         mlab.show()

    def forward(self, batch_dict):
        """
        Args:
            batch_dict:
                batch_size:
                point_features: (N1 + N2 + N3 + ..., C) or (B, N, C)
                point_coords: (N1 + N2 + N3 + ..., 4) [bs_idx, x, y, z]
                point_labels (optional): (N1 + N2 + N3 + ...)
                gt_boxes (optional): (B, M, 8)
        Returns:
            batch_dict:
                point_cls_scores: (N1 + N2 + N3 + ..., 1)
                point_part_offset: (N1 + N2 + N3 + ..., 3)
        """
        point_features = batch_dict['point_features']
        point_coords = batch_dict['point_coords']
        batch_size = batch_dict['batch_size']
        encoded_features = batch_dict['spatial_features_2d']
        encoded_tensor_stride = batch_dict['spatial_features_stride']
        encoded_features_to_pointwise = self.interpolate_from_bev_features(
            keypoints=point_coords,
            bev_features=encoded_features,
            batch_size=batch_size,
            bev_stride=encoded_tensor_stride,
        )
        pos_encode = self.pos_encode_layers(point_coords[:, 1:])
        encoded_features_to_pointwise = self.encoded_feature_layers(encoded_features_to_pointwise + pos_encode)
        # encoded_features_to_pointwise = self.encoded_feature_layers(encoded_features_to_pointwise)
        point_features = self.point_feature_layers(point_features)

        semantic_geo_features_base = torch.cat([encoded_features_to_pointwise,
                                                point_features, ], dim=1)
        encoded_features_to_pointwise_attention = torch.sigmoid(self.ef_attention_layers(semantic_geo_features_base))
        point_features_attention = torch.sigmoid(self.pf_attention_layers(semantic_geo_features_base))
        encoded_features_to_pointwise = encoded_features_to_pointwise*encoded_features_to_pointwise_attention
        point_features = point_features*point_features_attention
        semantic_geo_features = torch.cat([encoded_features_to_pointwise,
                                                point_features, ], dim=1)
        # semantic_geo_features = self.semantic_geo_layers(semantic_geo_features)

        point_cls_preds = self.cls_layers(semantic_geo_features)  # (total_points, num_class)
        point_part_preds = self.part_reg_layers(semantic_geo_features)

        ret_dict = {
            'point_cls_preds': point_cls_preds,
            'point_part_preds': point_part_preds,
        }
        if self.box_layers is not None:
            point_box_preds = self.box_layers(semantic_geo_features)
            ret_dict['point_box_preds'] = point_box_preds

        # _, point_coords_wlhr = self.generate_predicted_boxes(
        #     points=batch_dict['point_coords'][:, 1:4],
        #     point_cls_preds=point_cls_preds, point_box_preds=ret_dict['point_box_preds']
        # )

        # end_point_coords = torch.cat([point_coords[:, 0:1], point_coords_wlhr[:, 0:3]],
        #                              dim=1).detach()
        # batch_dict['end_point_coords'] = end_point_coords

        point_cls_scores = torch.sigmoid(point_cls_preds)
        point_part_offset = torch.sigmoid(point_part_preds)
        # score_for_show = point_part_offset.mean(dim=1)*point_cls_scores.max(dim=-1)[0]
        # self.show_point_score(
        #     batch_size=2,
        #     point_coords=point_coords,
        #     cls_scores=score_for_show
        # )
        batch_dict['point_cls_scores'], _ = point_cls_scores.max(dim=-1)
        batch_dict['point_part_offset'] = point_part_offset
        batch_dict['multi_scale_3d_features']['x_points'].features = semantic_geo_features

        if self.training:
            targets_dict = self.assign_targets(batch_dict)
            ret_dict['point_cls_labels'] = targets_dict['point_cls_labels']
            ret_dict['point_part_labels'] = targets_dict.get('point_part_labels')
            ret_dict['point_box_labels'] = targets_dict.get('point_box_labels')

        if self.box_layers is not None and (not self.training or self.predict_boxes_when_training):
            point_cls_preds, point_box_preds = self.generate_predicted_boxes(
                points=batch_dict['point_coords'][:, 1:4],
                point_cls_preds=point_cls_preds, point_box_preds=ret_dict['point_box_preds']
            )
            batch_dict['batch_cls_preds'] = point_cls_preds
            batch_dict['batch_box_preds'] = point_box_preds
            batch_dict['batch_index'] = batch_dict['point_coords'][:, 0]
            batch_dict['cls_preds_normalized'] = False

        self.forward_ret_dict = ret_dict
        return batch_dict


#!/usr/bin/env python

'''
python cluster.py --models ../../models/ --outdir ../../results --force --genomes /Volumes/container/genomic_files/

# --4vis vis.tsv
# better just include in output

TODO: do we need metadata? can we not just dereplicate and done? for plotting
for the publication we can code stuff by hand, like looking up clusters etc.:

> The clustering and dereplication based on `nanotext` is entirely
_taxonomy-free_ -- i.e. we need know nothing about these genomes and can
nevertheless dereplicate them confidently.
'''

import argparse
from collections import defaultdict
from glob import glob
import json
import os
from pathlib import Path

# Filter the multitude of warnings thrown by the HDBSCAN library
# stackoverflow.com/questions/879173
def warn(*args, **kwargs):
    pass
import warnings
warnings.warn = warn

import numpy as np
import pandas as pd
# from sklearn.manifold import TSNE
from tqdm import tqdm
import umap
import hdbscan
# from openTSNE import TSNE, TSNEEmbedding, affinity, initialization

from nanotext.classes import GenomeModel
from nanotext.io import dbopen
from nanotext.utils import get_taxa_from_names, get_names_from_taxon, cosine
from nanotext.utils import strip_name

parser = argparse.ArgumentParser(description='Process some integers.')

parser.add_argument(
    '--models', type=str, required=True,
    help='Path to models (directory)')

# parser.add_argument(
#     '--metadata', type=str, required=True,
#     help='Path to metadata SQLite database')

parser.add_argument(
    '--config', type=str, default='config.json',
    help='Path to config (file)')

parser.add_argument(
    '--outdir', type=str, required=True,
    help='Path to results (directory)')

parser.add_argument(
    '--genomes', type=str, required=True,
    help='Path to genomes (directory)')

# default=False is implied by action='store_true'
# https://stackoverflow.com/questions/8259001
parser.add_argument(
    '--reduce', action='store_true',
    help='Reduce dimensions before clustering')

parser.add_argument(
    '--force', action='store_true',
    help='Overwrite results if exist')

# TODO: rm? bc/ this is part of some bash script
# parser.add_argument(
#     '--min_identity', dest='ani', default=0.95, type=float,
#     help='ANI above which clustered genomes are dereplicated')

args = parser.parse_args()

in_ = Path(args.genomes).resolve(strict=True)
# TODO: this has to be the absolute path, not relative!
# https://stackoverflow.com/questions/42513056
# p.resolve(strict=True)
# returns a Path object
# Note: strict=True needs Python >= v3.6
# print(in_)

out = Path(args.outdir).resolve()  # does not exist yet, so no strict=True
if args.force:
    out.mkdir(exist_ok=True)
else:
    out.mkdir()

with open(args.config, 'r') as file:
    config = json.load(file)

model = GenomeModel(args.models, mode='core', norm='l2')
genomes = list(in_.glob('[!^\.]*.fasta'))
# Don't include hidden files (e.g. the OS index on a Mac)

print('Collecting embedding for each genome ...')
names, m = [], []
for g in genomes:
    name = g.name.replace('.fasta', '')
    try:
        v = model.embedding[name]
        names.append(name)
        m.append(v)
    except KeyError:
        continue
# names, m = zip(*[(k, v) for k, v in model.embedding.items()])
percent = int(100*len(m)/len(genomes))
print(f'Found {len(m)} embeddings for {len(genomes)} genomes ({percent}%)')

# For testing:
# m = m[:1000]
'''

'''

# on min_cluster_size and min_samples
# (3, 3) -- 3260 clusters, 29548 outliers
# (2, 2) -- 8744 clusters, 38053 outliers
# (3, 2) -- 6537 clusters, 35306 outliers

# (3, 3, 100 dim) -- 3800 clusters, 18k outliers
# (3, 1, 100 dim) -- 4344 clusters, 16109 outliers

'''
Calling clusters --------------------------------------------------------------

On reducing the dimensionality of the embedding before clustering:

> UMAP can be used as an effective preprocessing step to boost the performance
of density based clustering. This is somewhat controversial, and should be
attempted with care. For a good discussion of some of the issues involved in
this please see the various answers in this stackoverflow thread (1) on
clustering the results of t-SNE.
-- https://umap-learn.readthedocs.io/en/latest/clustering.html

(1) https://stats.stackexchange.com/questions/263539/clustering-on-the-output-of-t-sne

We used "enhanced clustering".
-- https://umap-learn.readthedocs.io/en/latest/clustering.html#umap-enhanced-clustering

>  [B]ut euclidean distance on l2 normalized vectors is equivalent to angular
distance, so you can normalize the vectors and then use euclidean distance.
-- https://github.com/scikit-learn-contrib/hdbscan/issues/69

https://hdbscan.readthedocs.io/en/latest/parameter_selection.html
https://hdbscan.readthedocs.io/en/latest/outlier_detection.html

Param selection:

"umap": {
    "random_state": 42,
    "n_components": 10,
    "metric": "cosine",
    "n_neighbors": 30,
    "min_dist": 0.5
}
Clusters: 5304
Outliers: 15266

- more mindist less clusters/ outliers

"hdbscan": {
    "min_cluster_size": 3,
    "min_samples": 1,
    "metric": "euclidean",
    "cluster_selection_method": "eom"
}
Clusters: 4344
Outliers: 16109
'''

clusterer = hdbscan.HDBSCAN(**config['hdbscan'])

if args.reduce:
    print('Reducing dimensions ...')
    redux = umap.UMAP(**config['umap']).fit_transform(m)
    # There is no random_state for the current HDBSCAN implementation
    print(f'Clustering {len(m)} points ...')
    cluster_labels = clusterer.fit_predict(redux)
else:
    print('No dimension reduction applied')
    print(f'Clustering {len(m)} points ...')
    cluster_labels = clusterer.fit_predict(m)

print(f'Clusters: {len(set(cluster_labels))}')
print(f'Outliers: {sum(cluster_labels == -1)}')

'''
Vis ---------------------------------------------------------------------------
'''

# print(m)

# TODO: optimize
print('Projecting points for visualization ...')
# projection = TSNE(n_components=2, metric='cosine', n_jobs=8).fit(np.array(m))
from sklearn.manifold import TSNE
projection = TSNE(n_components=2).fit_transform(np.array(m))

# TODO: Add vis? Simply uncomment.
l = []
outdirs = defaultdict(list)
for i, label in tqdm(enumerate(cluster_labels)):
    name = names[i]

    c1, c2 = projection[i]
    # c1, c2 = 0, 0  # hack until proper vis flag

    # Collect coords for visualization
    l.append([name, label, c1, c2])
    # We will dereplicate the clusters in parallel, so each cluster needs
    # to get a separate directory in the results.
    outdirs[label].append(name)

coords = pd.DataFrame.from_records(l)
coords.columns = 'acc label c1 c2'.split()
coords.to_csv(out/'4vis.csv', header=True, index=False)

'''
Create symlinks

https://kite.com/python/examples/1484/os-create-a-symbolic-link
'''

outc = out/'clusters'
outc.mkdir(exist_ok=True)  # will clobber if exists

print('Adding softlinks ...')
for label, names in outdirs.items():
    # Create a folder for each cluster (outliers are labelled -1).
    # When using this script in nextflow, it wants to feed a dir "-1" to
    # the dereplicator.py script, which is wrongly interpreted as a command
    # line argument -- rename.
    if label == -1:
        label = "outliers"

    fp = outc/str(label)
    fp.mkdir(exist_ok=True)
    # Add symbolic links to input genomes to each cluster directory
    for name in names:
        genome = f'{name}.fasta'
        target = fp/genome
        target.symlink_to(in_/genome)
'''
We will make use of the scripts from Ryan Wick to create the necessary files
to compute a Centrifuge index. For this, we need the GTDB taxonomy and the
genomes meant to go into the index. The genome IDs in the GTDB (taxonomy) are
prefixed w/ "GB_GCA..." (Genbank) and "RS_GCF" (Refseq). However, the index
scripts already prune those names, and so the filenames of the genomes need to
be formatted like "GCA_xxx.fasta".
'''


#!/usr/bin/env python
# coding: utf-8

# # Lista 03 - Análise Exploratória de Dados
#
# Continuando da última lista, vamos agora realizar um pouco dos passos da análise exploratória de dados. Em particular, vamos passar pelos passos de:
#
# 1. Carregamento dos dados
# 1. Limpeza dos dados
# 1. Análise exploratória com gráficos e estatísticas simples
#
# ## Imports Básicos
#
# As células abaixo apenas configuram nosso notebook para ficar mais parecido com os das aulas

# In[1]:

from numpy.testing import assert_almost_equal
from numpy.testing import assert_equal

from numpy.testing import assert_array_almost_equal
from numpy.testing import assert_array_equal

# In[2]:

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

# In[3]:

plt.rcParams['figure.figsize']  = (16, 10)
plt.rcParams['axes.labelsize']  = 20
plt.rcParams['axes.titlesize']  = 20
plt.rcParams['legend.fontsize'] = 20
plt.rcParams['xtick.labelsize'] = 20
plt.rcParams['ytick.labelsize'] = 20
plt.rcParams['lines.linewidth'] = 4

# In[4]:

plt.ion()
plt.style.use('seaborn-colorblind')

# ## Notas dos Alunos (Tutorial)
#
# Como falamos na última lista, em muitos cenários do mundo real, os dados são carregados de fontes como arquivos. Vamos substituir o DataFrame das notas dos alunos pelo conteúdo de um arquivo de texto.

# In[5]:

df = pd.read_csv('https://raw.githubusercontent.com/icd-ufmg/icd-ufmg.github.io/master/listas/l3/grades.csv',
                 sep=',',
                 header=0)
df.head()

# O método `read_csv` do `DataFrame` é usado para carregar dados de arquivos de texto. Como você pode ver no código de exemplo, você pode especificar opções como o delimitador de coluna e qual linha (se houver) contém cabeçalhos de coluna (neste caso, o delimitador é uma vírgula e a primeira linha contém os nomes das colunas).
#
# Além do mais, a chamada `head` imprime as primeiras cinco linhas da nossa tabela.

# ### Dados Faltantes
#
# Um dos problemas mais comuns com os quais os cientistas de dados precisam lidar são dados incompletos ou ausentes. Como podemos saber que o DataFrame contém valores ausentes? Você pode usar o método `isnull` para tal tarefa.

# In[6]:

df.isnull()

# Obseve como a última linha falta com o número de horas estudadas. Nas dúas últimas, faltam as notas. Caso você deseja saber a quantidade de dados faltantes, basta somar os `True`s da tabela acima. Aqui é importante saber que Python tratta `True` de forma similar ao número 1. Portanto, basta você somar a tabela inteira para pegar tal quantidade de dados faltantes.

# In[7]:

df.isnull().sum()

# Lembrando que a chamada `iloc` pega uma linha com base no número da mesma, vamos observar a última linha da tabela de dados. Observe como os valores faltantes viram `NaN`s. Além do mais, lembre-se que podemos indexar de trás para frente com o -1. -1 é a última linha, -2 a penúltima. Para entender a lógica, em um vetor de tamanho `n`, `n-1` é o último elemento. Indexar `-1` indica `n-1`.

# In[8]:

df.iloc[-1]

# A penúltima.

# In[9]:

df.iloc[-2]

# Agora que encontramos os valores faltantes, o que podemos fazer a respeito deles?
#
# #### fillna
#
# Uma abordagem comum é imputar valores de substituição. Por exemplo, se o número de horas de estudo está faltando, podemos simplesmente supor que o aluno estudou por um período médio de tempo e substituir o valor faltante com as horas de estudo médias. Para fazer isso, podemos usar o método fillna, como este:

# In[10]:

df['StudyHours'].fillna(df['StudyHours'].mean())

# Observe que a última linha foi alterada! Porém, o DataFrame original não foi.

# In[11]:

df.iloc[-1]

# Para alterar, podemos trocar a coluna. Abaixo faço tal operação em uma cópia dos dados. Realizei tal escolha apenas para não mudar a tabela original.

# In[12]:

df_novo = df.copy() # criar uma cópia apenas para o exemplo
df_novo['StudyHours'] = df_novo['StudyHours'].fillna(df_novo['StudyHours'].mean())
df_novo.iloc[-1]

# Observe como não mudamos nada das notas. O `fillna` pode receber uma série indexada para alterar várias colunas. Primeiramente, observe como a chamada `mean` pega a média de todas as colunas.

# In[13]:

df.mean()

# O `fillna` então vai pegar o índice dessa série, o nome da coluna, e utilizar como chave para quais colunas imputar. O valor da série indica o valor que será imputado. Observe como os novos dados abaixo estão sem NaNs.

# In[14]:

df_novo = df.fillna(df.mean())
df_novo

# In[15]:

df_novo.isnull().sum()

# #### dropna
#
# Outra opção é simplesmente remover todas as linhas com dados faltantes. Para tal, fazemos uso da chamada `dropna`.

# In[16]:

df_novo = df.dropna()
df_novo.shape

# In[17]:

df.shape

# Observe como o novo DataFrame tem duas linhas a menos do que o anterior. A escolha de como limpar dados faltantes depende do tipo de análise que você vai realizar. Aqui, vamos seguir com o drop no `df`.

# In[18]:

df.dropna(inplace=True) # on inplace=True altera o dataframe atual, não retorna um novo
df

# ### Explorando Dados
#
# Lembre-se que você pode indexar DataFrames com vetores booleanos. Por exemplo, para pegas as notas de Skye podemos primeiramente achar a discente nos dados:

# In[19]:

df['Name'] == 'Skye'

# Obserne como acima temos uma entrada verdadeira. Esta, é justamente a linha onde temos o nome `Skye`. Ao indexar o DataFrame com tal linha, pegamos a nota da discente.

# In[20]:

idx = df['Name'] == 'Skye'
df[idx]

# Outra forma de fazer a mesma operação é com o método **query**. O query faz consultados usando uma línguagem similar aos bancos de dados. Um exemplo:

# In[21]:

df.query('Name == "Skye"')

# Podemos também pegar todas as notas acima de 60. Ou seja, os alunos aprovados.

# In[22]:

df.query('Grade >= 60')

# Como também os alunos que passaram estudando relativamente pouco.  No nosso caso, vamos focar em alunos que estudaram menos do que 14horas.

# In[23]:

df.query('Grade >= 60 and StudyHours <= 14')

# Todo retorno, ou do índice booleano ou da query são outros DataFrames. Então, podemos chamar métodos como tirar a média dos alunos.

# In[24]:

above_60_low_hours = df.query('Grade >= 60 and StudyHours <= 14')
type(above_60_low_hours)

# In[25]:

above_60_low_hours.mean()

# Podemos também buscar os alunos que estão acima da média!

# In[26]:

mean = df['Grade'].mean()
df[df['Grade'] >= mean]

# ou, via query.

# In[27]:

df.query(f'Grade >= {mean}')

# Os DataFrames são incrivelmente versáteis e facilitam a manipulação de dados. Muitas operações DataFrame retornam uma nova cópia do DataFrame; portanto, se quiser modificar um DataFrame, mas manter a variável existente, você precisará atribuir o resultado da operação à variável existente. Por exemplo, o código a seguir classifica os dados do aluno em ordem decrescente de nota e atribui o DataFrame classificado resultante à variável `df_students`.

# In[28]:

# Re-ordena os dados por nota
df_students = df.sort_values('Grade', ascending=False)
df_students

# ### Visualizando dados com Matplotlib
#
# Os DataFrames fornecem uma ótima maneira de explorar e analisar dados tabulares, mas uma imagem vale mil palavras. A biblioteca [Matplotlib](matplotlib.org) fornece a base para a plotagem de visualizações de dados.
#
# Vamos começar com um histograma de notas. Observe como também colocamos uma linha preta em cada barra `edgecolor='k'` e setamos rótulos ao X e Y (para sabermos qual eixo mostra quais dados).

# In[29]:

plt.hist(df_students['Grade'], edgecolor='k')
plt.xlabel('Grade')
plt.ylabel('Num. Students')

# Observe como as notas concentram em 50. Às vezes é mais simples interpretar a função cumulativa dos dados.  Vamos fazer este gráfico de uma forma diferente da que vimos em sala de aula. Alterando um pouco a chamada `hist` podemos pegar a cumulativa.

# In[30]:

plt.hist(df_students['Grade'], edgecolor='k', cumulative=True)
plt.xlabel('Grade - x')
plt.ylabel('Num. Students with Grade <= x')

# Agora, `density=True` transforma este plot em uma função cumulativa de probabilidade.

# In[31]:

plt.hist(df_students['Grade'], edgecolor='k', cumulative=True, density=True)
plt.xlabel('Grade - x')
plt.ylabel('Frac. Students with Grade <= x')

# Observe que 50% dos alunos (0.5 no eixo-y) tem nota menor ou igual à mais ou menos 47~50 (eixo-x). Isto pode ser verificado com a chamada median abaixo. Outro exemplo, tente entender no gráfico, 20% dos alunos (eixo-y) tem nota menor ou igual à mais ou menos 30.

# In[32]:

df_students['Grade'].median()

# Até agora, você usou métodos do Matplotlib.pyplot para plotar gráficos. No entanto, muitos pacotes, incluindo Pandas, fornecem métodos que abstraem as funções Matplotlib simplificando sua vida. Por exemplo, o DataFrame fornece seus próprios métodos para plotar dados, conforme mostrado no exemplo a seguir para plotar um gráfico de barras de horas de estudo.

# In[33]:

df_students.plot.bar(x='Name', y='Grade', edgecolor='k')

# Ou o mesmo histograma de antes.

# In[34]:

df_students.plot.hist(y='Grade', edgecolor='k')
plt.xlabel('Grade - x')
plt.ylabel('Num. Students')

# Como também uma versão contínua do histograma. Esta é chamada de Kernel Density Estimation (vimos rapidamente em sala de aula).

# In[35]:

df_students.plot.kde(y='Grade')
plt.xlabel('Grade - x')
plt.ylabel('Density')

# ### Estatísticas
#
# #### Medidas de tendência central
#
# Para entender melhor a distribuição, podemos examinar as chamadas medidas de tendência central; que é uma maneira sofisticada de descrever estatísticas que representam o "meio" dos dados. O objetivo disso é tentar encontrar um valor "típico". Maneiras comuns de definir o meio dos dados incluem:
#   * A média: uma média simples baseada na soma de todos os valores no conjunto de amostra e, em seguida, na divisão do total pelo número de amostras.
#   * A mediana: o valor no meio do intervalo de todos os valores de amostra.
#   * A moda: o valor de ocorrência mais comum no conjunto de amostra.
#
# Vamos calcular esses valores, junto com os valores mínimo e máximo para comparação, e mostrá-los no histograma. Primeiramente, observe como podemos pegar tais valores direto do DataFrame.

# In[36]:

df.mean()

# In[37]:

df.median()

# Ou, descrever o DataFrame como um todo. Aqui pegamos a média, mínimo, máximo e quartis (mais abaixo) dos dados.

# In[38]:

df.describe()

# Além disso, podemos focar em um vetor de dados específico.

# In[39]:

data = df['Grade']

min_ = data.min()
max_ = data.max()
mean = data.mean()
median = data.median()
mode = data.mode()[0]

print(min_, max_, mean, median, mode, sep=', ')

# Abaixo plotamos cada estatística em linhas verticais.

# In[40]:

# Histograma
plt.hist(data, edgecolor='k')
plt.xlabel('Grade')
plt.ylabel('Num. Students')

# Linhas para cada estatística
plt.axvline(x=min_, color = 'gray', linestyle='dashed', linewidth = 2, label='min')
plt.axvline(x=mean, color = 'cyan', linestyle='dashed', linewidth = 2, label='mean')
plt.axvline(x=median, color = 'red', linestyle='dashed', linewidth = 2, label='median')
plt.axvline(x=mode, color = 'yellow', linestyle='dashed', linewidth = 2, label='mode')
plt.axvline(x=max_, color = 'gray', linestyle='dashed', linewidth = 2, label='max')

# Adiciona uma legenda
plt.legend()

# ### Medidas de variância
#
# Portanto, agora temos uma boa ideia de onde estão as estatísticas centrais dos dados. No entanto, há outro aspecto das distribuições que devemos examinar: quanta variabilidade existe nos dados?
#
# As estatísticas típicas que medem a variabilidade nos dados incluem:
#
# * Intervalo: a diferença entre o máximo e o mínimo. Não há função incorporada para isso, mas é fácil calcular usando as funções mín e máx.
# * Variância: a média da diferença quadrática da média. Você pode usar a função var integrada para encontrar isso.
# * Desvio padrão: a raiz quadrada da variância. Você pode usar a função std embutida para encontrar isso.

# In[41]:

for col_name in ['Grade', 'StudyHours']:
    col = df_students[col_name]
    rng = col.max() - col.min()
    var = col.var(ddof=1)
    std = col.std(ddof=1)
    print('\n{}:\n - Range: {:.2f}\n - Variance: {:.2f}\n - Std.Dev: {:.2f}'.format(col_name, rng, var, std))

# #### Quartis
#
# Na estatística descritiva, um quartil é qualquer um dos três valores que divide o conjunto ordenado de dados em quatro partes iguais, e assim cada parte representa 1/4 da amostra ou população.
#
# Assim, no caso duma amostra ordenada,
#   * primeiro quartil (designado por Q1/4) = quartil inferior = é o valor aos 25% da amostra ordenada = 25º percentil
#   * segundo quartil (designado por Q2/4) = mediana = é o valor até ao qual se encontra 50% da amostra ordenada = 50º percentil, ou 5º decil.
#   * terceiro quartil (designado por Q3/4) = quartil superior = valor a partir do qual se encontram 25% dos valores mais elevados = valor aos 75% da amostra ordenada = 75º percentil
#   * à diferença entre os quartis superior e inferior chama-se amplitude inter-quartil.
#
# Observe como os quartis estão presentes na sumarização dos dados. São as linhas 25%, 50% (mediana) e 75%.

# In[42]:

df.describe()

# #### Comparando Dados
#
# Por fim, vamos comparar o tempo de estudo entre os alunos que passaram ou não. Para tal, vamos definir qualquer nota >= 60 como sendo uma aprovação. Depois disso, vamos alterar o DataFrame para conter tal informação.
#
# Primeiro criando a série com tal informação.

# In[43]:

passed = df['Grade'] >= 60
passed

# Agora, alterando o DataFrame.

# In[44]:

df['Passed'] = passed
df

# Agora, vamos entender o tempo de dedicação dos alunos que passam e dos que não passam. Podemos fazer uso do Boxplot para tal comparativo. Na estatística descritiva, o boxplot é uma ferramenta gráfica para representar a variação de dados observados de uma variável numérica por meio de quartis.
#
# A "caixa" mostra o primeiro e terceiro quartil. A linha no meio mostra o segundo, a mediana. Portanto, a caixa se estende dos valores dos quartil de Q1 a Q3 dos dados, com uma linha na mediana (Q2). Os bigodes se estendem das bordas da caixa para mostrar a extensão dos dados. Por padrão, eles estendem não mais do que 1,5 * IQR (IQR = Q3 - Q1) das bordas da caixa, terminando no ponto de dados mais distante dentro desse intervalo. Valores fora desta faix são plotados como pontos separados.

# In[45]:

df.boxplot(column='StudyHours', by='Passed', grid=False)

# Observe que, como esperado, alunos que passam estudam mais. Os quartis, portanto a distribuição dos dados, são mais altos. Para observar os valores podemos fazer um groupby.
#
# Abaixo eu agrupo os dados por passou ou não. Depois, pego as horas de estudo. Por fim, observo as estatíticas.

# In[46]:

df.groupby('Passed')['StudyHours'].describe()

# Observe como 25% (primeiro quartil) dos alunos que não passam, estudam menos do que 8,25 horas. Para os que passam, este valor é de 13,125. Agora pense no complemento dos 25%. 75% dos alunos que passam estudam pelo menos 13h! Isso é bem maior do que as 8h dos que não passam!

# ### Correlacionando Dados
#
# Por fim, podemos correlacional o tempo de estudo com a nota.

# In[47]:

plt.scatter(df.StudyHours, df.Grade, edgecolor='k', alpha=0.75)
plt.xlabel('Hours')
plt.ylabel('Grade')

# ## Exercícios (Flight Data)
#
# Eu espero que o tutorial acima tenha sido uma boa revisão de análise exploratória de dados. Agora é com você! Nesta lista, você explorará um conjunto de dados do mundo real contendo dados de voos do Departamento de Transporte dos EUA.
#
# Vamos começar carregando e visualizando os dados.

# In[48]:

df = pd.read_csv('https://raw.githubusercontent.com/icd-ufmg/icd-ufmg.github.io/master/listas/l3/flights.csv')

# O conjunto de dados contém observações de voos domésticos dos EUA em 2013 e consiste nos seguintes campos:
# * Ano: o ano do voo (todos os registros são de 2013)
# * Mês: o mês do voo
# * Dia do mês: o dia do mês em que o voo partiu
# * DayOfWeek: o dia da semana em que o voo partiu - de 1 (segunda-feira) a 7 (domingo)
# * Transportadora: a abreviatura de duas letras da companhia aérea.
# * OriginAirportID: Um identificador numérico exclusivo para o aeroporto de partida
# * Nome do aeroporto de origem: o nome completo do aeroporto de partida
# * OriginCity: a cidade do aeroporto de partida
# * Estado de origem: o estado do aeroporto de partida
# * DestAirportID: Um identificador numérico único para o aeroporto de destino
# * DestAirportName: o nome completo do aeroporto de destino
# * DestCity: a cidade do aeroporto de destino
# * DestState: o estado do aeroporto de destino
# * CRSDepTime: a hora de partida programada
# * DepDelay: o número de minutos de atraso na partida (o voo que saiu antes do horário tem um valor negativo)
# * DelDelay15: Um indicador binário de que a partida foi atrasada por mais de 15 minutos (e, portanto, considerada "atrasada")
# * CRSArrTime: a hora de chegada programada
# * ArrDelay: o número de minutos de atraso na chegada (o voo que chegou antes do horário tem um valor negativo)
# * ArrDelay15: Um indicador binário de que a chegada foi atrasada em mais de 15 minutos (e, portanto, considerada "atrasada")
# * Cancelado: um indicador binário de que o voo foi cancelado

# In[49]:

df.head()

# ### Exercício 1
#
# Conte a quantidade de dados faltantes na tabela. Isto é, em TODAS as células. O método retorna apenas um número.

# In[50]:

def count_missing(df) -> int:
    return df.isna().sum().sum()

# In[51]:

assert_equal(2761, count_missing(df))

# ### Exercício 2
#
# Crie um novo DataFrame sem as linhas com dados faltantes

# In[52]:

def drop_missing(df):
    return df.dropna()

# In[53]:

assert_equal(269179, drop_missing(df).shape[0])

# ### Exercício 3
#
# Retorne a mediana de TODAS as colunas numéricas do DataFrame

# In[54]:

def all_median(df):
    return df.median()

# In[55]:

assert_equal(2013, all_median(df)['Year'])
assert_equal(7, all_median(df)['Month'])

# ### Exercício 4
#
# Quais são os atrasos médios (médios) de partida e chegada? Retorne uma tupla.

# In[56]:

df.head()

# In[57]:

def delay(df) -> tuple:
    return (df['DepDelay'].mean(), df['ArrDelay'].mean())

# In[58]:

assert_almost_equal(10.350448628373906, delay(df)[0])
assert_almost_equal(6.496649996322718, delay(df)[1])

# ### Exercício 5
#
# Indique qual rota tem o maior tempo de voo em MÉDIA. Use a coluna OriginAirportName e DestinationAirportName. Retorne uma tupla `(OriginAirportName, DestinationAirportName)`. Lembre-se de não considerar voos cancelados!

# In[59]:

def high_delay(df):
    df['Delay'] = df['CRSArrTime'] - df['CRSDepTime']
    df['Delay'] = df['Delay'].count()
    df = df.sort_values(by=['Delay'])
    return (df.iloc[0]['OriginAirportName'], df.iloc[0]['DestAirportName'])

# In[60]:

high_delay(df)

# ### Exercício 6
#
# Faça um boxplot dos atrasos de saída por dia da semana!
#
# 1. Dica, use `grid=False, showfliers=False` para o plot ficar mais limpo
#
# *Saída esperada*
#
# ![](https://raw.githubusercontent.com/icd-ufmg/icd-ufmg.github.io/master/listas/l3/plot1.png)

# In[63]:

df.boxplot(column='DepDelay', by='DayOfWeek', grid=False, showfliers=False)

# ### Exercício 7
#
# Correlacione o atraso de saída com o atraso de chegada!
#
# 1. Dica, remova os voos cancelados
#
# *Saída esperada*
#
# ![](https://raw.githubusercontent.com/icd-ufmg/icd-ufmg.github.io/master/listas/l3/plot2.png)

# In[69]:

df.loc[df['Cancelled'] == 0].plot.scatter(x='DepDelay', y='ArrDelay', )

# In[ ]:


from styx_msgs.msg import TrafficLight
import rospy

import tensorflow as tf
import os
import numpy as np
from collections import defaultdict
# from io import StringIO

PATH_TO_CKPT_REAL = './light_classification/models/ssd_inception_v2_real/frozen_inference_graph.pb'
PATH_TO_CKPT_SIM = './light_classification/models/ssd_inception_v2_sim/frozen_inference_graph.pb'

CATEGORY_INDEX = defaultdict(int, {1: TrafficLight.RED, 2: TrafficLight.YELLOW, 3: TrafficLight.GREEN})

CONFIDENCE_THRESHOLD = 0.6

class TLClassifier(object):
    def __init__(self, is_site):
        # Pick model based on current map
        if is_site:
        	model_path = PATH_TO_CKPT_REAL
	        rospy.loginfo("TLCLASSIFIER: Using Testing Site Model")
        else:
        	model_path = PATH_TO_CKPT_SIM
	        rospy.loginfo("TLCLASSIFIER: Using Simulation Model")

	    # Load model
        self.detection_graph = tf.Graph()
        with self.detection_graph.as_default():
			od_graph_def = tf.GraphDef()
			with tf.gfile.GFile(model_path, 'rb') as fid:
				serialized_graph = fid.read()
				od_graph_def.ParseFromString(serialized_graph)
				tf.import_graph_def(od_graph_def, name='')

		# Get identifiers from model graph
        with self.detection_graph.as_default():
			self.image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0')
			self.detection_boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0')
			self.detection_scores = self.detection_graph.get_tensor_by_name('detection_scores:0')
			self.detection_classes = self.detection_graph.get_tensor_by_name('detection_classes:0')
			self.num_detections = self.detection_graph.get_tensor_by_name('num_detections:0')

        self.sess = tf.Session(graph=self.detection_graph)

    def get_classification(self, image):

    	# Add an additional dimension to satisfy model's input requirement
        image_expanded = np.expand_dims(image, axis=0)

        # Run model
        (boxes, scores, classes, num) = self.sess.run(
		  [self.detection_boxes, self.detection_scores, self.detection_classes, self.num_detections],
		  feed_dict={self.image_tensor: image_expanded})

        # rospy.loginfo('class: {} \nscore: {}\n\n'.format(classes, scores))

        # Only cares about the best prediction (given that it passes the threshold)
        if (scores[0, 0] >= CONFIDENCE_THRESHOLD):
			light_pred = CATEGORY_INDEX[classes[0, 0]]
			return light_pred
        else:
			return TrafficLight.UNKNOWN


# -*- coding: utf-8 -*-

# Global imports
import argparse as ap
import numpy as np
from sklearn import cluster

# Local variable
Backbone_atoms = ["_CA_","_C__","_N__","_O__","_OXT","_HA_","_HA2","_HA3","_H__","_H1_","_H2_","_H3_"]

# Script information
__author__ = "Sergi Rodà"
__license__ = "MIT"
__version__ = "1.0.1"
__maintainer__ = "Sergi Rodà"
__email__ = "sergi.rodallordes@bsc.es"

def parseArgs():
	"""
	Parse arguments from command-line

	RETURNS
	-------
	traj : string
					list of trajectory files to analyze
	reference : string
					Path to the PDB reference file
	n_processors : integer
					number of processors to make the analysis
	cluster_width : float
					cluster width used in MeanShift algorithm
	non_ligand : boolean
					the boolean to decide where the ligand is clusterized or not
	residue_number : list of strings
					the residue number of the clusterized protein residues
	only_sidechain : boolean
					clusterize only the side chain atoms of the residues if true
	output : string
					filename of the output file
	"""

	parser = ap.ArgumentParser(description='Script that returns the density of the clusters of the ligand in a PELE trajectory')
	optional = parser._action_groups.pop()
	required = parser.add_argument_group('required arguments')
	parser.add_argument("traj", metavar="FILE",
							type=str, help="path to trajectory file")
	required.add_argument("-R", "--reference", required=True, metavar="FILE",
							type=str, help="path to PDB reference file")
	optional.add_argument("-n", "--n_processors", metavar="INTEGER",
							type=int, help="number of processors to make the analysis", default = None)
	optional.add_argument("-CW", "--cluster_width", metavar="FLOAT",
							type=float, help="cluster width in $\AA$", default = 1.5)
	optional.add_argument("-NL","--non_ligand",
							help = "Clusterize against a residue of the protein instead of the ligand", action = "store_true")
	optional.add_argument("-RN", "--residue_number", metavar="INTEGER",
							type=str, nargs='*', help="the residue numbers of the clusterized protein residues",default = [])
	optional.add_argument("-OS","--only_sidechain",
							help = "Clusterize only the atoms of the side chain of the residue/s", action = "store_true")
	optional.add_argument("-o","--output", metavar="PATH", type=str,help="filename of the output file", default="centroid")
	parser._action_groups.append(optional)
	args = parser.parse_args()

	return args.traj, args.reference, args.n_processors, args.cluster_width, args.non_ligand, args.residue_number, args.only_sidechain, args.output

def GetCoordinates(file, non_ligand, residue_number, only_sidechain, reference = False):
	"""
	This function returns the coordinates of a PDB file.

	PARAMETERS
	----------
	file : string
			PDB filename
	non_ligand : boolean
					the boolean to decide where the ligand is clusterized or not
	residue_number : list of strings
					the residue number of the clusterized protein residues
	only_sidechain : boolean

	RETURNS
	------
	res_coord: list of floats
	"""

	PDB = open(file)
	res_coord, res_coord_aux = [], []
	for line in PDB:
		if line[0:5]=="MODEL":
			if len(res_coord_aux) == 0:
				pass
			else:
				res_coord.append(list(np.mean(res_coord_aux,axis=0)))
				res_coord_aux = []
		if (line[21:22] == "L") and (non_ligand == False):
			x = float(line[30:38].strip())
			y = float(line[38:46].strip())
			z = float(line[46:54].strip())
			res_coord_aux.append([x,y,z])

		elif (non_ligand == True) and (line[22:26].strip() in residue_number) and (only_sidechain == False):
			x = float(line[30:38].strip())
			y = float(line[38:46].strip())
			z = float(line[46:54].strip())
			res_coord_aux.append([x,y,z])

		elif (non_ligand == True) and (line[22:26].strip() in residue_number) and (only_sidechain == True) and not (line[12:16].replace(" ","_") in Backbone_atoms):
			x = float(line[30:38].strip())
			y = float(line[38:46].strip())
			z = float(line[46:54].strip())
			res_coord_aux.append([x,y,z])

	if reference:
		res_coord = res_coord_aux

	return res_coord

def ClusterizeAtoms(traj, reference, non_ligand, residue_number, cluster_width, n_processors, only_sidechain, output):
	"""
	This function takes the arguments from the command line and clusterizes the position of the ligand
	(or the selected protein residues) during the PELE simulation according to the reference structure.

	PARAMETERS
	----------
	All command line arguments

	OUTPUT
	------
	A PDB file with the centroids of the average positions of the atom during the trajectory
	and the density is stored in the B-factor.
	"""

	# Retrieve trajectory data for cluster analysis
	Coords = GetCoordinates(traj, non_ligand, residue_number, only_sidechain)

	# Retrieve reference data for cluster analysis
	Ref_coords = GetCoordinates(reference, non_ligand, residue_number, only_sidechain, True)

	# Clustering
	Estimator = cluster.MeanShift(bandwidth = cluster_width, n_jobs = n_processors, cluster_all = True)
	Results = Estimator.fit_predict(Coords)

	Ref_clusters = []
	for ref_coord in Ref_coords:
				Ref_clusters += Estimator.predict([ref_coord]).tolist()

	# Clustering analysis
	aux = np.bincount(Results)
	aux_count = np.nonzero(aux)[0]
	density = {}
	for i in aux_count:
		density[i]=aux[i]/np.sum(aux)

	# Write centroids to PDB
	centroids = Estimator.cluster_centers_
	n_clusters = len(centroids)
	with open(output+".pdb", 'w') as f:
		for i, centroid in enumerate(centroids):
			f.write("ATOM    {:3d}  CEN BOX A {:3d} {:>11.3f}{:>8.3f}{:>8.3f}  1.00{:>5.2f}\n".format(i+1, i+1, *centroid, density[i]))

def main():
	"""
	Main function

	It is called when this script is the main program called by the interpreter
	"""

	traj, reference, n_processors, cluster_width, non_ligand, residue_number, only_sidechain, output = parseArgs()

	ClusterizeAtoms(traj, reference, non_ligand, residue_number, cluster_width, n_processors, only_sidechain, output)

if __name__=="__main__":
	main()