import os

import numpy as np
from paddle import fluid

from ltr.models.bbreg.atom import atom_resnet50, atom_resnet18
from ltr.models.siamese.siam import siamfc_alexnet
from ltr.models.siam.siam import SiamRPN_AlexNet, SiamMask_ResNet50_sharp, SiamMask_ResNet50_base
from pytracking.admin.environment import env_settings
from pytracking.features.featurebase import MultiFeatureBase
from pytracking.libs import TensorList
from pytracking.libs.paddle_utils import n2p


class ResNet18(MultiFeatureBase):
    """ResNet18 feature.
    args:
        output_layers: List of layers to output.
        net_path: Relative or absolute net path (default should be fine).
        use_gpu: Use GPU or CPU.
    """

    def __init__(self,
                 output_layers=('block2', ),
                 net_path='atom_iou',
                 use_gpu=True,
                 *args,
                 **kwargs):
        super().__init__(*args, **kwargs)

        self.output_layers = list(output_layers)
        self.use_gpu = use_gpu
        self.net_path = net_path

    def initialize(self):
        with fluid.dygraph.guard():
            if os.path.isabs(self.net_path):
                net_path_full = self.net_path
            else:
                net_path_full = os.path.join(env_settings().network_path,
                                             self.net_path)

            self.net = atom_resnet18(
                backbone_pretrained=False,
                backbone_is_test=True,
                iounet_is_test=True)

            state_dictsm, _ = fluid.load_dygraph(net_path_full)
            self.net.load_dict(state_dictsm)
            self.net.train()

            self.iou_predictor = self.net.bb_regressor

        self.layer_stride = {
            'conv0': 2,
            'conv1': 2,
            'block0': 4,
            'block1': 8,
            'block2': 16,
            'block3': 32,
            'classification': 16,
            'fc': None
        }
        self.layer_dim = {
            'conv0': 64,
            'conv1': 64,
            'block0': 64,
            'block1': 128,
            'block2': 256,
            'block3': 512,
            'classification': 256,
            'fc': None
        }

        self.iounet_feature_layers = self.net.bb_regressor_layer

        if isinstance(self.pool_stride, int) and self.pool_stride == 1:
            self.pool_stride = [1] * len(self.output_layers)

        self.feature_layers = sorted(
            list(set(self.output_layers + self.iounet_feature_layers)))

        self.mean = np.reshape([0.485, 0.456, 0.406], [1, -1, 1, 1])
        self.std = np.reshape([0.229, 0.224, 0.225], [1, -1, 1, 1])

    def free_memory(self):
        if hasattr(self, 'net'):
            del self.net
        if hasattr(self, 'iou_predictor'):
            del self.iou_predictor
        if hasattr(self, 'iounet_backbone_features'):
            del self.iounet_backbone_features
        if hasattr(self, 'iounet_features'):
            del self.iounet_features

    def dim(self):
        return TensorList([self.layer_dim[l] for l in self.output_layers])

    def stride(self):
        return TensorList([
            s * self.layer_stride[l]
            for l, s in zip(self.output_layers, self.pool_stride)
        ])

    def extract(self, im: np.ndarray, debug_save_name=None):
        with fluid.dygraph.guard():
            if debug_save_name is not None:
                np.savez(debug_save_name, im)

            im = im / 255.  # don't use im /= 255. since we don't want to alter the input
            im -= self.mean
            im /= self.std
            im = n2p(im)

            output_features = self.net.extract_features(im, self.feature_layers)

            # Store the raw resnet features which are input to iounet
            iounet_backbone_features = TensorList([
                output_features[layer] for layer in self.iounet_feature_layers
            ])
            self.iounet_backbone_features = iounet_backbone_features.numpy()

            # Store the processed features from iounet, just before pooling
            self.iounet_features = TensorList([
                f.numpy()
                for f in self.iou_predictor.get_iou_feat(
                    iounet_backbone_features)
            ])

            output = TensorList([
                output_features[layer].numpy() for layer in self.output_layers
            ])
            return output


class ResNet50(MultiFeatureBase):
    """ResNet50 feature.
    args:
        output_layers: List of layers to output.
        net_path: Relative or absolute net path (default should be fine).
        use_gpu: Use GPU or CPU.
    """

    def __init__(self,
                 output_layers=('block2', ),
                 net_path='atom_iou',
                 use_gpu=True,
                 *args,
                 **kwargs):
        super().__init__(*args, **kwargs)

        self.output_layers = list(output_layers)
        self.use_gpu = use_gpu
        self.net_path = net_path

    def initialize(self):
        with fluid.dygraph.guard():
            if os.path.isabs(self.net_path):
                net_path_full = self.net_path
            else:
                net_path_full = os.path.join(env_settings().network_path,
                                             self.net_path)

            self.net = atom_resnet50(
                backbone_pretrained=False,
                backbone_is_test=True,
                iounet_is_test=True)

            state_dictsm, _ = fluid.load_dygraph(net_path_full)
            self.net.load_dict(state_dictsm)
            self.net.train()

            self.iou_predictor = self.net.bb_regressor

        self.layer_stride = {
            'conv0': 2,
            'conv1': 2,
            'block0': 4,
            'block1': 8,
            'block2': 16,
            'block3': 32,
            'classification': 16,
            'fc': None
        }
        self.layer_dim = {
            'conv0': 64,
            'conv1': 64,
            'block0': 256,
            'block1': 512,
            'block2': 1024,
            'block3': 2048,
            'classification': 256,
            'fc': None
        }

        self.iounet_feature_layers = self.net.bb_regressor_layer

        if isinstance(self.pool_stride, int) and self.pool_stride == 1:
            self.pool_stride = [1] * len(self.output_layers)

        self.feature_layers = sorted(
            list(set(self.output_layers + self.iounet_feature_layers)))

        self.mean = np.reshape([0.485, 0.456, 0.406], [1, -1, 1, 1])
        self.std = np.reshape([0.229, 0.224, 0.225], [1, -1, 1, 1])

    def free_memory(self):
        if hasattr(self, 'net'):
            del self.net
        if hasattr(self, 'iou_predictor'):
            del self.iou_predictor
        if hasattr(self, 'iounet_backbone_features'):
            del self.iounet_backbone_features
        if hasattr(self, 'iounet_features'):
            del self.iounet_features

    def dim(self):
        return TensorList([self.layer_dim[l] for l in self.output_layers])

    def stride(self):
        return TensorList([
            s * self.layer_stride[l]
            for l, s in zip(self.output_layers, self.pool_stride)
        ])

    def extract(self, im: np.ndarray, debug_save_name=None):
        with fluid.dygraph.guard():
            if debug_save_name is not None:
                np.savez(debug_save_name, im)

            im = im / 255.  # don't use im /= 255. since we don't want to alter the input
            im -= self.mean
            im /= self.std
            im = n2p(im)

            output_features = self.net.extract_features(im, self.feature_layers)

            # Store the raw resnet features which are input to iounet
            iounet_backbone_features = TensorList([
                output_features[layer] for layer in self.iounet_feature_layers
            ])
            self.iounet_backbone_features = iounet_backbone_features.numpy()

            # Store the processed features from iounet, just before pooling
            self.iounet_features = TensorList([
                f.numpy()
                for f in self.iou_predictor.get_iou_feat(
                    iounet_backbone_features)
            ])

            output = TensorList([
                output_features[layer].numpy() for layer in self.output_layers
            ])
            return output


class SFCAlexnet(MultiFeatureBase):
    """Alexnet feature.
    args:
        output_layers: List of layers to output.
        net_path: Relative or absolute net path (default should be fine).
        use_gpu: Use GPU or CPU.
    """

    def __init__(self,
                 output_layers=('conv5', ),
                 net_path='estimator',
                 use_gpu=True,
                 *args,
                 **kwargs):
        super().__init__(*args, **kwargs)

        self.output_layers = list(output_layers)
        self.use_gpu = use_gpu
        self.net_path = net_path

    def initialize(self):
        with fluid.dygraph.guard():
            if os.path.isabs(self.net_path):
                net_path_full = self.net_path
            else:
                net_path_full = os.path.join(env_settings().network_path,
                                             self.net_path)

            self.net = siamfc_alexnet(
                backbone_pretrained=False,
                backbone_is_test=True,
                estimator_is_test=True)

            state_dictsm, _ = fluid.load_dygraph(net_path_full)
            self.net.load_dict(state_dictsm)
            self.net.train()

            self.target_estimator = self.net.target_estimator

        self.layer_stride = {'conv5': 8}
        self.layer_dim = {'conv5': 256}

        self.estimator_feature_layers = self.net.target_estimator_layer

        if isinstance(self.pool_stride, int) and self.pool_stride == 1:
            self.pool_stride = [1] * len(self.output_layers)

        self.feature_layers = sorted(
            list(set(self.output_layers + self.estimator_feature_layers)))

        self.mean = np.reshape([0., 0., 0.], [1, -1, 1, 1])
        self.std = np.reshape([1 / 255., 1 / 255., 1 / 255.], [1, -1, 1, 1])

    def free_memory(self):
        if hasattr(self, 'net'):
            del self.net
        if hasattr(self, 'target_estimator'):
            del self.target_estimator
        if hasattr(self, 'estimator_backbone_features'):
            del self.estimator_backbone_features

    def dim(self):
        return TensorList([self.layer_dim[l] for l in self.output_layers])

    def stride(self):
        return TensorList([
            s * self.layer_stride[l]
            for l, s in zip(self.output_layers, self.pool_stride)
        ])

    def extract(self, im: np.ndarray, debug_save_name=None):
        with fluid.dygraph.guard():
            if debug_save_name is not None:
                np.savez(debug_save_name, im)

            im = im / 255.  # don't use im /= 255. since we don't want to alter the input
            im -= self.mean
            im /= self.std
            im = n2p(im)

            output_features = self.net.extract_features(im, self.feature_layers)

            # Store the raw backbone features which are input to estimator
            estimator_backbone_features = TensorList([
                output_features[layer]
                for layer in self.estimator_feature_layers
            ])
            self.estimator_backbone_features = estimator_backbone_features.numpy(
            )

            output = TensorList([
                output_features[layer].numpy() for layer in self.output_layers
            ])
            return output


class SRPNAlexNet(MultiFeatureBase):
    """Alexnet feature.
    args:
        output_layers: List of layers to output.
        net_path: Relative or absolute net path (default should be fine).
        use_gpu: Use GPU or CPU.
    """

    def __init__(self,
                 net_path='estimator',
                 use_gpu=True,
                 *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.use_gpu = use_gpu
        self.net_path = net_path

    def initialize(self):
        with fluid.dygraph.guard():
            if os.path.isabs(self.net_path):
                net_path_full = self.net_path
            else:
                net_path_full = os.path.join(env_settings().network_path, self.net_path)

            self.net = SiamRPN_AlexNet(backbone_pretrained=False, is_test=True)

            state_dict, _ = fluid.load_dygraph(net_path_full)
            self.net.load_dict(state_dict)
            self.net.eval()

    def free_memory(self):
        if hasattr(self, 'net'):
            del self.net

    def extract(self, im: np.ndarray, debug_save_name=None):
        with fluid.dygraph.guard():
            if debug_save_name is not None:
                np.savez(debug_save_name, im)

            im = n2p(im)

            output_features = self.net.extract_backbone_features(im)

            # Store the raw backbone features which are input to estimator
            output = TensorList([layer.numpy() for layer in output_features])
            return output


class SMaskResNet50_base(MultiFeatureBase):
    """Resnet50-dilated feature.
    args:
        output_layers: List of layers to output.
        net_path: Relative or absolute net path (default should be fine).
        use_gpu: Use GPU or CPU.
    """

    def __init__(self,
                 net_path='estimator',
                 use_gpu=True,
                 *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.use_gpu = use_gpu
        self.net_path = net_path

    def initialize(self):
        with fluid.dygraph.guard():
            if os.path.isabs(self.net_path):
                net_path_full = self.net_path
            else:
                net_path_full = os.path.join(env_settings().network_path, self.net_path)

            self.net = SiamMask_ResNet50_base(backbone_pretrained=False, is_test=True)

            state_dict, _ = fluid.load_dygraph(net_path_full)
            self.net.load_dict(state_dict)
            self.net.eval()

    def free_memory(self):
        if hasattr(self, 'net'):
            del self.net

    def extract(self, im: np.ndarray, debug_save_name=None):
        with fluid.dygraph.guard():
            if debug_save_name is not None:
                np.savez(debug_save_name, im)

            im = n2p(im)

            output_features = self.net.extract_backbone_features(im)

            # Store the raw backbone features which are input to estimator
            output = TensorList([layer.numpy() for layer in output_features])
            return output


class SMaskResNet50_sharp(MultiFeatureBase):
    """Resnet50-dilated feature.
    args:
        output_layers: List of layers to output.
        net_path: Relative or absolute net path (default should be fine).
        use_gpu: Use GPU or CPU.
    """

    def __init__(self,
                 net_path='estimator',
                 use_gpu=True,
                 *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.use_gpu = use_gpu
        self.net_path = net_path

    def initialize(self):
        with fluid.dygraph.guard():
            if os.path.isabs(self.net_path):
                net_path_full = self.net_path
            else:
                net_path_full = os.path.join(env_settings().network_path, self.net_path)

            self.net = SiamMask_ResNet50_sharp(backbone_pretrained=False, is_test=True)

            state_dict, _ = fluid.load_dygraph(net_path_full)
            self.net.load_dict(state_dict)
            self.net.eval()

    def free_memory(self):
        if hasattr(self, 'net'):
            del self.net

    def extract(self, im: np.ndarray, debug_save_name=None):
        with fluid.dygraph.guard():
            if debug_save_name is not None:
                np.savez(debug_save_name, im)

            im = n2p(im)

            output_features = self.net.extract_backbone_features(im)

            # Store the raw backbone features which are input to estimator
            output = TensorList([layer.numpy() for layer in output_features])
            return output

[(83, 'numpy.reshape', 'np.reshape', (['[0.485, 0.456, 0.406]', '[1, -1, 1, 1]'], {}), True, 'import numpy as np\n'), (84, 'numpy.reshape', 'np.reshape', (['[0.229, 0.224, 0.225]', '[1, -1, 1, 1]'], {}), True, 'import numpy as np\n'), (97, 'pytracking.libs.TensorList', 'TensorList', (['[self.layer_dim[l] for l in self.output_layers]'], {}), False, 'from pytracking.libs import TensorList\n'), (204, 'numpy.reshape', 'np.reshape', (['[0.485, 0.456, 0.406]', '[1, -1, 1, 1]'], {}), True, 'import numpy as np\n'), (205, 'numpy.reshape', 'np.reshape', (['[0.229, 0.224, 0.225]', '[1, -1, 1, 1]'], {}), True, 'import numpy as np\n'), (218, 'pytracking.libs.TensorList', 'TensorList', (['[self.layer_dim[l] for l in self.output_layers]'], {}), False, 'from pytracking.libs import TensorList\n'), (307, 'numpy.reshape', 'np.reshape', (['[0.0, 0.0, 0.0]', '[1, -1, 1, 1]'], {}), True, 'import numpy as np\n'), (308, 'numpy.reshape', 'np.reshape', (['[1 / 255.0, 1 / 255.0, 1 / 255.0]', '[1, -1, 1, 1]'], {}), True, 'import numpy as np\n'), (319, 'pytracking.libs.TensorList', 'TensorList', (['[self.layer_dim[l] for l in self.output_layers]'], {}), False, 'from pytracking.libs import TensorList\n'), (36, 'paddle.fluid.dygraph.guard', 'fluid.dygraph.guard', ([], {}), False, 'from paddle import fluid\n'), (37, 'os.path.isabs', 'os.path.isabs', (['self.net_path'], {}), False, 'import os\n'), (43, 'ltr.models.bbreg.atom.atom_resnet18', 'atom_resnet18', ([], {'backbone_pretrained': '(False)', 'backbone_is_test': '(True)', 'iounet_is_test': '(True)'}), False, 'from ltr.models.bbreg.atom import atom_resnet50, atom_resnet18\n'), (48, 'paddle.fluid.load_dygraph', 'fluid.load_dygraph', (['net_path_full'], {}), False, 'from paddle import fluid\n'), (106, 'paddle.fluid.dygraph.guard', 'fluid.dygraph.guard', ([], {}), False, 'from paddle import fluid\n'), (113, 'pytracking.libs.paddle_utils.n2p', 'n2p', (['im'], {}), False, 'from pytracking.libs.paddle_utils import n2p\n'), (118, 'pytracking.libs.TensorList', 'TensorList', (['[output_features[layer] for layer in self.iounet_feature_layers]'], {}), False, 'from pytracking.libs import TensorList\n'), (157, 'paddle.fluid.dygraph.guard', 'fluid.dygraph.guard', ([], {}), False, 'from paddle import fluid\n'), (158, 'os.path.isabs', 'os.path.isabs', (['self.net_path'], {}), False, 'import os\n'), (164, 'ltr.models.bbreg.atom.atom_resnet50', 'atom_resnet50', ([], {'backbone_pretrained': '(False)', 'backbone_is_test': '(True)', 'iounet_is_test': '(True)'}), False, 'from ltr.models.bbreg.atom import atom_resnet50, atom_resnet18\n'), (169, 'paddle.fluid.load_dygraph', 'fluid.load_dygraph', (['net_path_full'], {}), False, 'from paddle import fluid\n'), (227, 'paddle.fluid.dygraph.guard', 'fluid.dygraph.guard', ([], {}), False, 'from paddle import fluid\n'), (234, 'pytracking.libs.paddle_utils.n2p', 'n2p', (['im'], {}), False, 'from pytracking.libs.paddle_utils import n2p\n'), (239, 'pytracking.libs.TensorList', 'TensorList', (['[output_features[layer] for layer in self.iounet_feature_layers]'], {}), False, 'from pytracking.libs import TensorList\n'), (278, 'paddle.fluid.dygraph.guard', 'fluid.dygraph.guard', ([], {}), False, 'from paddle import fluid\n'), (279, 'os.path.isabs', 'os.path.isabs', (['self.net_path'], {}), False, 'import os\n'), (285, 'ltr.models.siamese.siam.siamfc_alexnet', 'siamfc_alexnet', ([], {'backbone_pretrained': '(False)', 'backbone_is_test': '(True)', 'estimator_is_test': '(True)'}), False, 'from ltr.models.siamese.siam import siamfc_alexnet\n'), (290, 'paddle.fluid.load_dygraph', 'fluid.load_dygraph', (['net_path_full'], {}), False, 'from paddle import fluid\n'), (328, 'paddle.fluid.dygraph.guard', 'fluid.dygraph.guard', ([], {}), False, 'from paddle import fluid\n'), (335, 'pytracking.libs.paddle_utils.n2p', 'n2p', (['im'], {}), False, 'from pytracking.libs.paddle_utils import n2p\n'), (340, 'pytracking.libs.TensorList', 'TensorList', (['[output_features[layer] for layer in self.estimator_feature_layers]'], {}), False, 'from pytracking.libs import TensorList\n'), (371, 'paddle.fluid.dygraph.guard', 'fluid.dygraph.guard', ([], {}), False, 'from paddle import fluid\n'), (372, 'os.path.isabs', 'os.path.isabs', (['self.net_path'], {}), False, 'import os\n'), (377, 'ltr.models.siam.siam.SiamRPN_AlexNet', 'SiamRPN_AlexNet', ([], {'backbone_pretrained': '(False)', 'is_test': '(True)'}), False, 'from ltr.models.siam.siam import SiamRPN_AlexNet, SiamMask_ResNet50_sharp, SiamMask_ResNet50_base\n'), (379, 'paddle.fluid.load_dygraph', 'fluid.load_dygraph', (['net_path_full'], {}), False, 'from paddle import fluid\n'), (388, 'paddle.fluid.dygraph.guard', 'fluid.dygraph.guard', ([], {}), False, 'from paddle import fluid\n'), (392, 'pytracking.libs.paddle_utils.n2p', 'n2p', (['im'], {}), False, 'from pytracking.libs.paddle_utils import n2p\n'), (419, 'paddle.fluid.dygraph.guard', 'fluid.dygraph.guard', ([], {}), False, 'from paddle import fluid\n'), (420, 'os.path.isabs', 'os.path.isabs', (['self.net_path'], {}), False, 'import os\n'), (425, 'ltr.models.siam.siam.SiamMask_ResNet50_base', 'SiamMask_ResNet50_base', ([], {'backbone_pretrained': '(False)', 'is_test': '(True)'}), False, 'from ltr.models.siam.siam import SiamRPN_AlexNet, SiamMask_ResNet50_sharp, SiamMask_ResNet50_base\n'), (427, 'paddle.fluid.load_dygraph', 'fluid.load_dygraph', (['net_path_full'], {}), False, 'from paddle import fluid\n'), (436, 'paddle.fluid.dygraph.guard', 'fluid.dygraph.guard', ([], {}), False, 'from paddle import fluid\n'), (440, 'pytracking.libs.paddle_utils.n2p', 'n2p', (['im'], {}), False, 'from pytracking.libs.paddle_utils import n2p\n'), (467, 'paddle.fluid.dygraph.guard', 'fluid.dygraph.guard', ([], {}), False, 'from paddle import fluid\n'), (468, 'os.path.isabs', 'os.path.isabs', (['self.net_path'], {}), False, 'import os\n'), (473, 'ltr.models.siam.siam.SiamMask_ResNet50_sharp', 'SiamMask_ResNet50_sharp', ([], {'backbone_pretrained': '(False)', 'is_test': '(True)'}), False, 'from ltr.models.siam.siam import SiamRPN_AlexNet, SiamMask_ResNet50_sharp, SiamMask_ResNet50_base\n'), (475, 'paddle.fluid.load_dygraph', 'fluid.load_dygraph', (['net_path_full'], {}), False, 'from paddle import fluid\n'), (484, 'paddle.fluid.dygraph.guard', 'fluid.dygraph.guard', ([], {}), False, 'from paddle import fluid\n'), (488, 'pytracking.libs.paddle_utils.n2p', 'n2p', (['im'], {}), False, 'from pytracking.libs.paddle_utils import n2p\n'), (108, 'numpy.savez', 'np.savez', (['debug_save_name', 'im'], {}), True, 'import numpy as np\n'), (229, 'numpy.savez', 'np.savez', (['debug_save_name', 'im'], {}), True, 'import numpy as np\n'), (330, 'numpy.savez', 'np.savez', (['debug_save_name', 'im'], {}), True, 'import numpy as np\n'), (390, 'numpy.savez', 'np.savez', (['debug_save_name', 'im'], {}), True, 'import numpy as np\n'), (438, 'numpy.savez', 'np.savez', (['debug_save_name', 'im'], {}), True, 'import numpy as np\n'), (486, 'numpy.savez', 'np.savez', (['debug_save_name', 'im'], {}), True, 'import numpy as np\n'), (40, 'pytracking.admin.environment.env_settings', 'env_settings', ([], {}), False, 'from pytracking.admin.environment import env_settings\n'), (161, 'pytracking.admin.environment.env_settings', 'env_settings', ([], {}), False, 'from pytracking.admin.environment import env_settings\n'), (282, 'pytracking.admin.environment.env_settings', 'env_settings', ([], {}), False, 'from pytracking.admin.environment import env_settings\n'), (375, 'pytracking.admin.environment.env_settings', 'env_settings', ([], {}), False, 'from pytracking.admin.environment import env_settings\n'), (423, 'pytracking.admin.environment.env_settings', 'env_settings', ([], {}), False, 'from pytracking.admin.environment import env_settings\n'), (471, 'pytracking.admin.environment.env_settings', 'env_settings', ([], {}), False, 'from pytracking.admin.environment import env_settings\n')]