from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Dict, List

from design import DiagramNode, Entity, Input, Mechanic, Stock
from design.links import Link
from design.sheep_wolf_diagram import SHEEP_WOLF_DIAGRAM


def calculate_design_metrics() -> Dict[str, Any]:
    """
    Calculate basic structural metrics for the default IMR diagram.
    """

    diagram = SHEEP_WOLF_DIAGRAM

    nodes: List[DiagramNode] = list(diagram.nodes.values())
    links: List[Link] = list(diagram.links)

    input_nodes = [n for n in nodes if isinstance(n, Input)]
    entity_nodes = [n for n in nodes if isinstance(n, Entity)]
    stock_nodes = [n for n in nodes if isinstance(n, Stock)]
    mechanic_nodes = [n for n in nodes if isinstance(n, Mechanic)]

    # Inputs: count and sum of degrees of freedom (numeric values).
    num_inputs = len(input_nodes)
    sum_dof = sum(float(inp.degree_of_freedom) for inp in input_nodes)

    # Entities: total count.
    num_entities = len(entity_nodes)

    # Stocks: total count and sum of magnitude orders (numeric values).
    num_stocks = len(stock_nodes)
    sum_magnitude_orders = sum(float(stock.magnitude_order) for stock in stock_nodes)

    # Mechanics: total count.
    num_mechanics = len(mechanic_nodes)

    # Attributes: in this diagram, modeled as entities without stock.
    num_attributes = sum(1 for e in entity_nodes if e.stock is None)

    # Links: total count.
    num_links = len(links)

    return {
        "inputs": {
            "count": num_inputs,
            "sum_dof": sum_dof,
        },
        "entities": {
            "count": num_entities,
        },
        "stocks": {
            "count": num_stocks,
            "sum_magnitude_order": sum_magnitude_orders,
        },
        "mechanics": {
            "count": num_mechanics,
        },
        "attributes": {
            "count": num_attributes,
        },
        "links": {
            "count": num_links,
        },
    }


def main() -> None:
    metrics = calculate_design_metrics()

    output_path = Path(__file__).with_name("design_metrics.json")
    output_path.write_text(json.dumps(metrics, indent=2))

    print(f"Design metrics written to {output_path}")


if __name__ == "__main__":
    main()


