import copy
import hyperchamber as hc
import inspect
import math
import operator
import os
import re
from functools import reduce

import pyparsing
import hypergan
import torch
import torch.nn as nn

from .gan_component import GANComponent
from hypergan.gan_component import ValidationException
from hypergan.layer_shape import LayerShape
from hypergan.distributions.base_distribution import BaseDistribution

from hypergan.modules.adaptive_instance_norm import AdaptiveInstanceNorm
from hypergan.modules.attention import Attention
from hypergan.modules.const import Const
from hypergan.modules.learned_noise import LearnedNoise
from hypergan.modules.modulated_conv2d import ModulatedConv2d, Blur, EqualLinear
from hypergan.modules.multi_head_attention import MultiHeadAttention
from hypergan.modules.reshape import Reshape
from hypergan.modules.no_op import NoOp
from hypergan.modules.scaled_conv2d import ScaledConv2d
from hypergan.modules.variational import Variational
from hypergan.modules.pixel_norm import PixelNorm

import torchvision
import hypergan as hg

class ConfigurableComponent(GANComponent):
    custom_layers = {}
    def __init__(self, gan, config, input=None, input_shape=None, context_shapes = {}, input_is_latent=False):
        self.current_size = LayerShape(gan.channels(), gan.height(), gan.width())
        if isinstance(input, GANComponent):
            if hasattr(input, 'current_height'):
                self.current_size = LayerShape(input.current_channels, input.current_height, input.current_width)
            elif hasattr(input, 'current_channels'):
                self.current_size = LayerShape(input.current_channels)
            else:
                self.current_size = input.current_size
        if input_shape is not None:
            self.current_size = LayerShape(*input_shape)
        self.layers = []
        self.layer_shapes = []
        self.untrainable_parameters = set()
        self.layer_output_sizes = {}
        self.nn_layers = []
        self.layer_options = {}
        self.parsed_layers = []
        self.parser = hypergan.parser.Parser()
        self.context_shapes = context_shapes
        for key, shape in self.context_shapes.items():
            self.layer_output_sizes[key] = shape
        if isinstance(input, BaseDistribution):
            self.is_latent = True
        else:
            self.is_latent = False
        self._latent_parameters = []
        self.layer_ops = {**self.activations(),
            **ConfigurableComponent.custom_layers,
            "add": hg.layers.Add,
            "cat": hg.layers.Cat,
            "channel_attention": hg.layers.ChannelAttention,
            "efficient_attention": hg.layers.EfficientAttention,
            "ez_norm": hg.layers.EzNorm,
            "layer": hg.layers.Layer,
            "minibatch": hg.layers.Minibatch,
            "mul": hg.layers.Mul,
            "multi_head_attention2": hg.layers.MultiHeadAttention, #TODO rename
            "noise": hg.layers.Noise,
            "pixel_shuffle": hg.layers.PixelShuffle,
            "residual": hg.layers.Residual,
            "resizable_stack": hg.layers.ResizableStack,
            "segment_softmax": hg.layers.SegmentSoftmax,
            "upsample": hg.layers.Upsample,

            #easy to convert
            "dropout": self.layer_dropout,
            "identity": self.layer_identity,
            "flatten": self.layer_flatten,
            "pretrained": self.layer_pretrained,
            "avg_pool": self.layer_avg_pool,#TODO handle dims
            "pad": self.layer_pad,
            "reshape": self.layer_reshape,
            "split": self.layer_split,

            #hard to convert
            "adaptive_avg_pool": self.layer_adaptive_avg_pool,
            "adaptive_avg_pool1d": self.layer_adaptive_avg_pool1d,
            "adaptive_avg_pool3d": self.layer_adaptive_avg_pool3d,
            "adaptive_instance_norm": self.layer_adaptive_instance_norm,
            "attention": self.layer_attention,
            "batch_norm": self.layer_batch_norm,
            "batch_norm1d": self.layer_batch_norm1d,
            "blur": self.layer_blur,
            "const": self.layer_const,
            "conv": self.layer_conv,
            "conv1d": self.layer_conv1d,
            "conv2d": self.layer_conv2d,
            "conv3d": self.layer_conv3d,
            "deconv": self.layer_deconv,
            "equal_linear": self.layer_equal_linear,
            "instance_norm": self.layer_instance_norm,
            "instance_norm1d": self.layer_instance_norm1d,
            "instance_norm3d": self.layer_instance_norm3d,
            "latent": self.layer_latent,
            "layer_norm": self.layer_norm,
            "learned_noise": self.layer_learned_noise,
            "linear": self.layer_linear,
            "modulated_conv2d": self.layer_modulated_conv2d,
            "multi_head_attention": self.layer_multi_head_attention,
            "pixel_norm": self.layer_pixel_norm,
            "resize_conv": self.layer_resize_conv,
            "resize_conv2d": self.layer_resize_conv2d,
            "resize_conv1d": self.layer_resize_conv1d,
            "scaled_conv2d": self.layer_scaled_conv2d,
            "subpixel": self.layer_subpixel,
            "vae": self.layer_vae
            #"linear_attention": hg.layers.LinearAttention,
            #"make2d": self.layer_make2d,
            #"make3d": self.layer_make3d,
            # "crop": self.layer_crop,
            # "dropout": self.layer_dropout,
            # "noise": self.layer_noise, #TODO
            #"attention": self.layer_attention, #TODO
            #"const": self.layer_const, #TODO
            #"gram_matrix": self.layer_gram_matrix, #TODO
            #"image_statistics": self.layer_image_statistics, #TODO
            #"knowledge_base": self.layer_knowledge_base, #TODO
            #"layer_norm": self.layer_layer_norm,#TODO
            #"mask": self.layer_mask,#TODO
            #"match_support": self.layer_match_support,#TODO
            #"pixel_norm": self.layer_pixel_norm,#TODO
            #"progressive_replace": self.layer_progressive_replace,#TODO
            #"reduce_sum": self.layer_r
educe_sum,#TODO might want to just do "reduce sum" instead
            #"relational": self.layer_relational,#TODO
            #"unpool": self.layer_unpool, #TODO https:/
/arxiv.org/abs/1505.04366
            #"squash": self.layer_squash, #TODO
            #"tensorflowcv": self.layer_tensorflowcv, #TODO layer torchvision instead?
            #"turing_test": self.layer_turing_test, #TODO
            #"two_sample_stack": self.layer_two_sample_stack, #TODO
            #"zeros": self.layer_zeros, #TODO
            #"zeros_like": self.layer_zeros_like #TODO
            }
        self.named_layers = {}
        if not hasattr(gan, "named_layers"):
            gan.named_layers = {}
        self.subnets = hc.Config(hc.Config(config).subnets or {})
        GANComponent.__init__(self, gan, config)
        self.device = self.config.device or "cuda:0"
        self.const_two = torch.Tensor([2.0]).float()[0].cuda()
        self.const_one = torch.Tensor([1.0]).float()[0].cuda()

    def required(self):
        return "layers".split()

    def layer(self, name):
        if name in self.gan.named_layers:
            return self.gan.named_layers[name] 
        if name in self.named_layers:
            return self.named_layers[name]
        return None

    def create(self):
        for layer in self.config.layers:
            net = self.create_parsed_layer(layer)
            self.nn_layers.append(net)

        self.net = nn.ModuleList(self.nn_layers)

    def create_parsed_layer(self, layer_defn):
        config = self.config
        parsed, layer = self.parse_layer(layer_defn)
        self.parsed_layers.append(parsed)
        self.layer_shapes.append(self.current_size)
        return layer

    def parse_layer(self, layer_defn):
        print("Parsing layer:", layer_defn)
        parsed = self.parser.parse_string(layer_defn)
        parsed.parsed_options = hc.Config(parsed.options)
        parsed.layer_defn = layer_defn
        print("Parsed layer:", parsed.to_list())
        layer = self.build_layer(parsed.layer_name, parsed.args, parsed