from collections.abc import Callable
from copy import copy, deepcopy
from functools import partial, singledispatch, singledispatchmethod
from typing import Any

import numpy as np
import tensorflow as tf
from keras.src.engine.keras_tensor import KerasTensor
from keras.src.engine.node import Node
from tensorflow import keras

from ..layers import HLayerBase, HQuantize, PDropout, PLayerBase, Signature
from ..layers.base import ABSBaseLayer
from ..utils import warn
from .fixed_point_quantizer import FixedPointQuantizer
from .precision_derivation import register_qconf
from .unary_lut import xfr_to_unary_lut


def get_all_nodes(model: keras.Model) -> set[Node]:
    """Get all nodes in the model as a set."""
    nodes = set()
    layers = set(model.layers)
    for layer in model.layers:
        for node in layer._inbound_nodes:
            if node.layer in layers:
                nodes.add(node)
        for node in layer._outbound_nodes:
            if node.layer in layers:
                nodes.add(node)
    return nodes


def solve_dependencies(model: keras.Model):
    """Given a keras model, return the input nodes, output nodes and a list of (layer, requires, provides) tuples. Requires is a list of nodes that are parents of the layer, provides is the node that is the output of the layer."""
    inp_tensors: list[KerasTensor] = model.inputs  # type:ignore
    out_tensors: list[KerasTensor] = model.outputs  # type:ignore

    input_nodes: list[Node] = [t.node for t in inp_tensors]  # type:ignore
    output_nodes: list[Node] = [t.node for t in out_tensors]  # type:ignore

    nodes = get_all_nodes(model)

    dependencies_dict: dict[Node, tuple[keras.layers.Layer, list[Node]]] = {}
    "List of (layer, requires, provides) tuples; requires is a list of nodes, provides is a single node."

    for node in nodes:
        if node.is_input:
            continue
        layer = node.layer  # layer that is called on the node
        requires = list(node.parent_nodes)
        provides = node
        dependencies_dict[provides] = (layer, requires)
    return input_nodes, output_nodes, dependencies_dict


def get_weight(layer: keras.layers.Layer, name: str):
    """Given a layer and a weight name, return the weight. The weight name may or may not contain the layer name. If the number index is missing, it is assumed to be 0."""
    if '/' in name:
        name = name.split('/')[-1]
    if ':' not in name:
        name = f'{name}:0'
    for w in layer.weights:
        if name in w.name:
            return w
    return None


def copy_fused_weights(src: keras.layers.Layer, dst: keras.layers.Layer):
    """For HGQ layers, some layers may have different fused weights for kernel and bias (Processed weights are deployment). This function copies the fused kernel and bias to the keras proxy."""
    if hasattr(dst, 'kernel') and dst.kernel is not None:
        if (k := getattr(src, 'fused_qkernel', None)) is not None:
            dst.kernel.assign(k)
        else:
            dst.kernel.assign(src.kernel)
    if hasattr(dst, 'bias') and dst.bias is not None:
        if (b := getattr(src, 'fused_qbias', None)) is not None:
            dst.bias.assign(b)
        else:
            dst.bias.assign(src.bias)
    for dst_w in dst.weights:
        if 'kernel:0' in dst_w.name or 'bias:0' in dst_w.name:
            continue
        src_w = get_weight(src, dst_w.name)
        if src_w is not None:
            dst_w.assign(src_w)
        else:
            warn(f'No weight found for layer {src.name} corresponding to {dst_w.name}.')


class Namer:
    """Helper class to generate unique names for layers, if one being used multiple times."""

    def __init__(self):
        self.used_names: set[str] = set()

    def next_name(self, name: str):
        i = 0
        if name in self.used_names:
            while f'{name}_{i}' in self.used_names:
                i += 1
            name = f'{name}_{i}'
        self.used_names.add(name)
        return name


def _get_out_tensor(
    dependencies_dict: dict[Node, tuple[keras.layers.Layer, list[Node]]],
    satisfied: dict[Node, Any],
    node: Node,
    namer: Namer | None = None,
    layer_xformer: Callable[[keras.layers.Layer], keras.layers.Layer] = lambda x: x,
):
    """Given a node, get the output tensor provided generated by the node."""

    assert node not in satisfied, f'Node {node} already satisfied.'

    namer = namer or Namer()
    layer, requires = dependencies_dict[node]
    inps = []
    for req in requires:
        if req in satisfied:
            inps.append(satisfied[req])
        else:
            inp = _get_out_tensor(dependencies_dict, satisfied, req, namer, layer_xformer)
            inps.append(inp)
    if len(inps) == 1:
        inps = inps[0]

    satisfied[node] = apply_layer(layer, inps, namer=namer, layer_xformer=layer_xformer)
    return satisfied[node]


@singledispatch
def apply_layer(
    layer: keras.layers.Layer | None,
    inp_tensors: list[tf.Tensor] | tf.Tensor,
    namer: None | Namer = None,
    # satisfied: dict[Node, Any] | None = None,
    layer_xformer: Callable[[keras.layers.Layer], keras.layers.Layer] = lambda x: x,
):
    """
    Apply a layer to the input tensors:
      - if the layer is a model, apply the model recursively in a flattened manner.
      - if the layer is a keras layer, transform it by layer_transformer, until the exactly same reference is returned.
    """

    layer_xf = layer_xformer(layer)
    n = 0
    namer = namer or Namer()
    while layer_xf is not layer:
        if isinstance(layer_xf, keras.Model):
            # layer_transformer may return a keras.Model from a layer.
            return apply_layer(layer_xf, inp_tensors, namer=namer, layer_xformer=layer_xformer)
        assert n < 1024, f'layer_transformer does not converge for layer (name: {layer.name if layer is not None else None}).'
        layer = layer_xf
        layer_xf = layer_xformer(layer)
        n += 1

    if layer_xf is not None:
        name = namer.next_name(layer_xf.name)
        layer_xf._name = name
        # Remove all inbound and outbound nodes to clean up the graph.
        layer_xf = copy(layer_xf)
        layer_xf._inbound_nodes = []
        layer_xf._outbound_nodes = []
        return layer_xf(inp_tensors)
    else:
        return inp_tensors


@apply_layer.register
def _(
    layer: keras.Model,
    inp_tensors: list[tf.Tensor] | tf.Tensor,
    namer: None | Namer = None,
    # satisfied: dict[Node, Any] | None = None,
    layer_xformer: Callable[[keras.layers.Layer], keras.layers.Layer] = lambda x: x,
):
    model = layer
    if not model.built:
        if isinstance(inp_tensors, (list, tuple)):
            model.build([t.shape for t in inp_tensors])
        else:
            model.build(inp_tensors.shape)

    namer = namer or Namer()
    inp_tensors = [inp_tensors] if not isinstance(inp_tensors, (list, tuple)) else inp_tensors
    input_nodes, output_nodes, dependencies_dict = solve_dependencies(model)
    satisfied = {node: tensor for node, tensor in zip(input_nodes, inp_tensors)}
    get_out_tensor = partial(_get_out_tensor, dependencies_dict, satisfied, namer=namer, layer_xformer=layer_xformer)
    outs = [get_out_tensor(output_node) for output_node in output_nodes]
    return outs[0] if len(outs) == 1 else outs


def convert_model(model: keras.Model, layer_xformer: Callable[[keras.layers.Layer], keras.layers.Layer] = lambda x: x):
    """For a keras model, convert each layer with layer_transformer, flatten all layers (remove all keras.Model as layers), and return a new model.

    Args:
        model (keras.Model): Input keras model.
        layer_transformer: A function that takes a keras layer and returns a keras layer (can be a keras.Model). For the final set of layer, it must return the same reference, rather than a copy, or the conversion will never end.

    Returns:
        keras.Model: A new keras model.

    Examples:
        Simply flatten a keras model: `convert_model(model)`
        Flatten a keras model, return a new model with all
    """

    inputs: list[tf.Tensor] = [keras.layers.Input(shape=t.shape[1:]) for t in model.inputs]  # type: ignore
    if len(inputs) == 1:
        inputs = inputs[0]
    outs = apply_layer(model, inputs, layer_xformer=layer_xformer)
    return keras.Model(inputs, outs)


@singledispatch
def to_keras_layer(layer):
    raise TypeError(f'No matching overload for layer type {type(layer)}. Signatures available: {to_keras_layer.registry.keys()}')


@to_keras_layer.register
def _(layer: ABSBaseLayer):
    """Given a HGQ layer, return the corresponding keras layer.

    Example:
        HDense -> Dense
        HQuantize -> Raise TypeError
    """

    assert not isinstance(layer, (HQuantize, Signature)), 'to_keras_layers should not be called on HQuantize or Signature layers.'

    if hasattr(layer, 'get_keras_config'):
        conf = layer.get_keras_config()
    else:
        conf = layer.get_config()

    if not isinstance(layer, keras.layers.Activation):
        if 'activation' in conf and conf['activation'] != 'linear':
            # Activation will be processed separately for non-activation layers.
            conf['activation'] = 'linear'
    else:
        # Prevent custom activation crashing conversion.
        conf['activation'] = layer.activation

    cls_name = layer.__class__.__name__[1:]
    if hasattr(keras.layers, cls_name):
        layer_cls = getattr(keras.layers, cls_name)
    elif hasattr(keras.layers, cls_name.replace('BatchNorm', '')):
        layer_cls = getattr(keras.layers, cls_name.replace('BatchNorm', ''))
    else:
        raise RuntimeError(f'Unknown layer type {layer.__class__.__name__}: no corresponding keras layer found.')
    klayer = layer_cls.from_config(conf)
    klayer.build(layer.input_shape)
    copy_fused_weights(layer, klayer)
    return klayer


class ProxyLayerXFormer:
    def __init__(self, SAT='WRAP', use_uniary_lut=False):
        self.SAT = SAT
        self.use_uniary_lut = use_uniary_lut

    @singledispatchmethod
    def get_kbiRS(self, layer: HLayerBase, **kwargs):
        q = layer.paq
        pos_only = kwargs.get('pos_only', False)
        k, i, f = q.get_bits_exact(pos_only=pos_only)
        k, b, i = k, k + i + f, k + i
        if q.rnd_strategy != 3 and not layer.can_bias_cover_rnd:
            RND = 'RND'
        else:
            RND = 'TRN'
        return k, b, i, RND, self.SAT

    @singledispatchmethod
    def __call__(self, layer: keras.layers.Layer | None) -> keras.layers.Layer | None:
        if layer is None:
            return None
        try:
            from hls4ml.converters.keras_to_hls import layer_handlers
        except ImportError:
            return layer
        cls_name = layer.__class__.__name__
        assert cls_name in layer_handlers, f'Layer {layer.name}: {cls_name} is not supported by hls4ml.'
        return layer

    @__call__.register
    def _(self, layer: HQuantize):
        kbiRS = self.get_kbiRS(layer)
        return FixedPointQuantizer(*kbiRS, name=layer.name)

    @__call__.register
    def _(self, layer: Signature):
        return FixedPointQuantizer(layer.keep_negative, layer.bits, layer.int_bits, 'TRN', self.SAT, name=layer.name)

    @__call__.register
    def _(self, layer: HLayerBase):
        # HLayer = Layer + Quantizer, or Layer + Quantizer1 + ReLU + Quantizer2

        name = layer.name
        klayer = to_keras_layer(layer)

        input_shape = layer.input_shape
        if isinstance(input_shape, tuple):
            inputs = keras.layers.Input(shape=input_shape[1:])
        else:
            inputs = [keras.layers.Input(shape=shape[1:]) for shape in input_shape]

        overrides = {'layers': {}}
        if hasattr(layer, 'parallel_factor') and layer.parallel_factor is not None:
            parallel_factor = layer.parallel_factor
            overrides = {'layers': {name: {'parallelization_factor': int(parallel_factor)}}}

        if not isinstance(klayer, keras.layers.Activation):
            if layer._relu_act:  # a non-activation layer can only have relu or liner as activation
                k, i, f = layer.paq.get_bits_exact(pos_only=False)
                k, i, f = tf.reduce_max(k), tf.reduce_max(i), tf.reduce_max(f)
                k, b, i = k, k + i + f, k + i
                rk, rb, ri, R, S = self.get_kbiRS(layer, pos_only=True)

                f_add_bits = 0 if R == 'TRN' else 1 if R == 'RND' else 2
                fq1 = FixedPointQuantizer(k, b + f_add_bits, i, SAT=S, RND='TRN', name=f'{name}_quantizer', overrides=overrides)
                fq2 = FixedPointQuantizer(rk, rb, ri, SAT=S, RND=R, name=f'{name}_relu_quantizer')
                return keras.Model(inputs, fq2(keras.layers.ReLU()(fq1(klayer(inputs)))))

        k, b, i, R, S = self.get_kbiRS(layer)
        q = FixedPointQuantizer(k, b, i, R, SAT=S, name=f'{name}_quantizer', overrides=overrides)
        return keras.Model(inputs, q(klayer(inputs)))

    @__call__.register
    def _(self, layer: PLayerBase):
        klayer = to_keras_layer(layer)
        return klayer

    @__call__.register
    def _(self, layer: keras.layers.Dropout):
        # Dropout does nothing at inference time
        return None

    @__call__.register
    def _(self, layer: keras.layers.Activation):
        # Purge linear activation on sight
        if layer.activation is keras.activations.linear:
            return None
        return layer


def to_proxy_model(model: keras.Model, aggressive: bool = True, accum_fp_max_offset: int | None = None, unary_lut_max_table_size=-1):
    """Given a HGQ model, return a hls4ml-ready keras model.

    Args:
        model: The HGQ model to be converted.

        aggressive (default: True): If True, use WRAP overflow mode. Significant performance degradation may occur if overflow occurs, but latency may be reduced. If False, use SAT overflow mode. Performance is more stable when it overflows, but latency may be increased.

        accum_fp_max_offset (default: None): If not None, autoset accumulator such that the model is bit accurate (when no overflow occurs and up to fp precision). If set, use the specified number of floating bits plus result float bits as accumulator float bits. May improve latency in some rare cases, not recommended in general.

        unary_lut_max_table_size (default: -1): If greater than 0, use unary LUT for `HActivation` layers, when the required table size is less than or equal to the specified value. If set to -1, do not use unary LUT.

    """

    if accum_fp_max_offset is not None and not aggressive:
        warn('You are using bitgrowth (aggressive=False) together with bias_accum_bits set. This is not recommended. If you are sure what you are doing, ignore this warning.')
    if accum_fp_max_offset is not None and accum_fp_max_offset < 0:
        warn('You are using a negative value for bias_accum_bits. Please make sure you know what you are doing.')
    proxy = convert_model(model, layer_xformer=ProxyLayerXFormer('WRAP' if aggressive else 'SAT').__call__)
    if unary_lut_max_table_size > 0:
        proxy = convert_model(proxy, layer_xformer=partial(xfr_to_unary_lut, max_table_size=unary_lut_max_table_size))
    for layer in proxy.layers:
        register_qconf(layer, accum_fp_max_offset=accum_fp_max_offset)
    return proxy
