# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

from typing import Iterable, Tuple, Optional
import numpy as np

import torch
from torch import nn, Tensor
from overrides import overrides

from archai.nas.cell import Cell
from archai.nas.operations import Op, DropPath_
from archai.nas.model_desc import ModelDesc, AuxTowerDesc, CellDesc
from archai.common import ml_utils
from archai.nas.arch_module import ArchModule

class Model(ArchModule):
    def __init__(self, model_desc:ModelDesc, droppath:bool, affine:bool):
        super().__init__()

        # some of these fields are public as finalizer needs access to them
        self.desc = model_desc

        # TODO: support any number of stems
        assert len(model_desc.model_stems)==2, "Model compiler currently only supports 2 stems"
        stem0_op = Op.create(model_desc.model_stems[0], affine=affine)
        stem1_op = Op.create(model_desc.model_stems[1], affine=affine)
        self.model_stems = nn.ModuleList((stem0_op, stem1_op))

        self.cells = nn.ModuleList()
        self._aux_towers = nn.ModuleList()

        for i, (cell_desc, aux_tower_desc) in \
                enumerate(zip(model_desc.cell_descs(), model_desc.aux_tower_descs)):
            self._build_cell(cell_desc, aux_tower_desc, droppath, affine)

        # adaptive pooling output size to 1x1
        self.pool_op = Op.create(model_desc.pool_op, affine=affine)
        # since ch_p records last cell's output channels
        # it indicates the input channel number
        self.logits_op = Op.create(model_desc.logits_op, affine=affine)

        # for i,cell in enumerate(self.cells):
        #     print(i, ml_utils.param_size(cell))
        #logger.info({'model_summary': self.summary()})

    def _build_cell(self, cell_desc:CellDesc,
                    aux_tower_desc:Optional[AuxTowerDesc],
                    droppath:bool, affine:bool)->None:
        trainables_from = None if cell_desc.trainables_from==cell_desc.id  \
                            else self.cells[cell_desc.trainables_from]
        cell = Cell(cell_desc, affine=affine, droppath=droppath,
                    trainables_from=trainables_from)
        self.cells.append(cell)
        self._aux_towers.append(AuxTower(aux_tower_desc) \
                                if aux_tower_desc else None)

    def summary(self)->dict:
        all_arch_params = list(self.all_owned()
                               .param_by_kind(kind=None))
        return {
            'cell_count': len(self.cells),
            #'cell_params': [ml_utils.param_size(c) for c in self.cells]
            'params': ml_utils.param_size(self),
            'arch_params_len': len(all_arch_params),
            'arch_params_numel': np.sum(a.numel() for a in all_arch_params),
            'ops': np.sum(len(n.edges) for c in self.desc.cell_descs() for n in c.nodes()),
        }

    def ops(self)->Iterable[Op]:
        for cell in self.cells:
            for op in cell.ops():
                yield op

    @overrides
    def forward(self, x)->Tuple[Tensor, Optional[Tensor]]:
        #print(torch.cuda.memory_allocated()/1.0e6)
        s0 = self.model_stems[0](x)
        #print(torch.cuda.memory_allocated()/1.0e6)
        s1 = self.model_stems[1](x)
        #print(-1, s0.shape, s1.shape, torch.cuda.memory_allocated()/1.0e6)

        logits_aux = None
        for ci, (cell, aux_tower) in enumerate(zip(self.cells, self._aux_towers)):
            #print(s0.shape, s1.shape, end='')
            s0, s1 = s1, cell.forward(s0, s1)
            #print(ci, s0.shape, s1.shape, torch.cuda.memory_allocated()/1.0e6)

            # TODO: this mimics darts but won't work for multiple aux towers
            if aux_tower is not None and self.training:
                logits_aux = aux_tower(s1)
                #print(ci, 'aux', logits_aux.shape)

        # s1 is now the last cell's output
        out = self.pool_op(s1)
        logits = self.logits_op(out) # flatten

        #print(-1, 'out', out.shape)
        #print(-1, 'logits', logits.shape)

        return logits, logits_aux

    def device_type(self)->str:
        return next(self.parameters()).device.type

    def drop_path_prob(self, p:float):
        """ Set drop path probability
        This will be called externally so any DropPath_ modules get
        new probability. Typically, every epoch we will reduce this probability.
        """
        for module in self.modules():
            if isinstance(module, DropPath_):
                module.p = p


class AuxTower(nn.Module):
    def __init__(self, aux_tower_desc:AuxTowerDesc):
        """assuming input size 14x14"""
        # TODO: assert input size?
        super().__init__()

        self.features = nn.Sequential(
            nn.ReLU(inplace=True),
            nn.AvgPool2d(5, stride=aux_tower_desc.stride, padding=0, count_include_pad=False),
            nn.Conv2d(aux_tower_desc.ch_in, 128, 1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 768, 2, bias=False),
            # TODO: This batchnorm was omitted in orginal implementation due to a typo.
            nn.BatchNorm2d(768),
            nn.ReLU(inplace=True),
        )
        self.logits_op = nn.Linear(768, aux_tower_desc.n_classes)

    def forward(self, x:torch.Tensor):
        x = self.features(x)
        x = self.logits_op(x.view(x.size(0), -1))
        return x