# Import necessary libraries
import gradio as gr
import sys
from huggingface_hub import ModelCard, HfApi
import requests
import networkx as nx
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from collections import defaultdict
from networkx.drawing.nx_pydot import graphviz_layout
from io import BytesIO
from PIL import Image

# Define the model ID
MODEL_ID = "mlabonne/NeuralBeagle14-7B"

# Define a class to cache model cards
class CachedModelCard(ModelCard):
  _cache = {}

  @classmethod
  def load(cls, model_id: str, **kwargs) -> "ModelCard":
    if model_id not in cls._cache:
      try:
        cls._cache[model_id] = super().load(model_id, **kwargs)
      except:
        cls._cache[model_id] = None
    return cls._cache[model_id]

# Function to get model names from a YAML file
def get_model_names_from_yaml(url):
    model_tags = []
    response = requests.get(url)
    if response.status_code == 200:
        model_tags.extend([item for item in response.content if '/' in str(item)])
    return model_tags

# Function to get the color of the model based on its license
def get_license_color(model):
    try:
        card = CachedModelCard.load(model)
        license = card.data.to_dict()['license'].lower()
        permissive_licenses = ['mit', 'bsd', 'apache-2.0', 'openrail']
        if any(perm_license in license for perm_license in permissive_licenses):
            return 'lightgreen'
        else:
            return 'lightcoral'
    except Exception as e:
        return 'lightgray'

# Function to find model names in the family tree
def get_model_names(model, genealogy, found_models=None, visited_models=None):
    if found_models is None:
        found_models = set()
    if visited_models is None:
        visited_models = set()

    if model in visited_models:
        return found_models
    visited_models.add(model)

    try:
        card = CachedModelCard.load(model)
        card_dict = card.data.to_dict()
        license = card_dict['license']

        model_tags = []
        if 'base_model' in card_dict:
            model_tags = card_dict['base_model']

        if 'tags' in card_dict and not model_tags:
            tags = card_dict['tags']
            model_tags = [model_name for model_name in tags if '/' in model_name]

        if not model_tags:
            model_tags.extend(get_model_names_from_yaml(f"https://huggingface.co/{model}/blob/main/merge.yml"))
        if not model_tags:
            model_tags.extend(get_model_names_from_yaml(f"https://huggingface.co/{model}/blob/main/mergekit_config.yml"))

        if not isinstance(model_tags, list):
            model_tags = [model_tags] if model_tags else []

        found_models.add(model)

        for model_tag in model_tags:
            genealogy[model_tag].append(model)
            get_model_names(model_tag, genealogy, found_models, visited_models)

    except Exception as e:
        pass

    return found_models

# Function to create the family tree
def create_family_tree(start_model):
    genealogy = defaultdict(list)
    get_model_names(start_model, genealogy)

    G = nx.DiGraph()

    for parent, children in genealogy.items():
        for child in children:
            G.add_edge(parent, child)

    max_depth = nx.dag_longest_path_length(G) + 1
    max_width = max_width_of_tree(G) + 1

    height = max(8, 1.6 * max_depth)
    width = max(8, 6 * max_width)

    plt.figure(figsize=(width, height))
    pos = graphviz_layout(G, prog="dot")

    node_colors = [get_license_color(node) for node in G.nodes()]
    clear_output()

    labels = {node: node.replace("/", "\n") for node in G.nodes()}

    nx.draw(G, pos, labels=labels, with_labels=True, node_color=node_colors, font_size=12, node_size=8_000, edge_color='black')

    legend_elements = [
        Patch(facecolor='lightgreen', label='Permissive'),
        Patch(facecolor='lightcoral', label='Noncommercial'),
        Patch(facecolor='lightgray', label='Unknown')
    ]
    plt.legend(handles=legend_elements, loc='upper left')

    plt.title(f"{start_model}'s Family Tree", fontsize=20)
    plt.show()

create_family_tree(MODEL_ID)