from typing import Tuple
from typing import Dict
from typing import List
from typing import Any


def clip_layers(clip_model) -> Tuple[Dict[str, int], List[Dict[str, Any]]]:
    """
    Gets names parameters in a structured way.
    Gives a tuple of {type_of_layer: count}, where type can be text, image, projection, tokens, or other
        and a list of dicts for each:
        {"type": str, "index": int, "param": torch_param, "name": orig_name}
            where type is as above, index is from 0 to count-1,
            and torch_param is the param as returned by module.named_parameters
    """
    classed_parameters = []
    metadata = {k: 0 for k in {"text", "image", "proj", "tokens", "other"}}
    for name, param in clip_model.named_parameters():
        # top layers always need to train
        if (
            name.startswith("ln_final.")
            or name.startswith("text_projection")
            or name.startswith("logit_scale")
            or name.startswith("visual.ln_post.")
            or name.startswith("visual.proj")
        ):
            t = "proj"
            inx = metadata[t]
        elif name.startswith("visual.transformer.resblocks."):
            t = "image"
            inx = int(name.split(".")[3])
        elif name.startswith("transformer.resblocks."):
            t = "text"
            inx = int(name.split(".")[2])
        elif name.startswith("token_embedding"):
            t = "tokens"
            inx = metadata[t]
        else:
            t = "other"
            inx = metadata[t]
        classed_parameters.append(
            {"type": t, "index": inx, "param": param, "name": name}
        )
        metadata[t] += 1
    for t in {"text", "image"}:
        metadata[t] = (
            max(
                classed_parameters, key=lambda cp: cp["index"] if cp["type"] == t else 0
            )["index"]
            + 1
        )

    return metadata, classed_parameters