import json

def series_availability(avails):
    result = 1.0
    for a in avails:
        result *= a
    return result

def parallel_availability(avails):
    result = 1.0
    for a in avails:
        result *= (1 - a)
    return 1 - result

def kofn_availability(avails, k):
    if k <= 1:
        return parallel_availability(avails)
    return sum(avails) / len(avails)  # simple approx

def compute_node_availability(node_id, schematic, schematics, overrides, visited):
    if node_id in visited:
        return 1.0
    visited.add(node_id)

    node = schematic["nodes"].get(node_id)
    if not node:
        return 1.0

    node_name = node.get("name", "").strip().lower()
    if node_name in overrides:
        return overrides[node_name]

    ntype = node["type"].lower()

    if ntype == "regularnode":
        return node.get("availability", 1.0)

    elif ntype == "series":
        return series_availability([
            compute_node_availability(child, schematic, schematics, overrides, visited.copy())
            for child in node.get("children", [])
        ])

    elif ntype == "parallel":
        return parallel_availability([
            compute_node_availability(child, schematic, schematics, overrides, visited.copy())
            for child in node.get("children", [])
        ])

    elif ntype == "kofn":
        k = node.get("min_required", 1)
        return kofn_availability([
            compute_node_availability(child, schematic, schematics, overrides, visited.copy())
            for child in node.get("children", [])
        ], k)

    elif ntype == "subschematic":
        tschematic_id = node.get("tschematic_id")
        if tschematic_id and str(tschematic_id) in schematics:
            sub_schematic = schematics[str(tschematic_id)]
            root = sub_schematic.get("root")
            if root:
                return compute_node_availability(root, sub_schematic, schematics, overrides, set())
        return node.get("availability", 1.0)

    return node.get("availability", 1.0)


def compute_schematic_availability(schematic, schematics, overrides):
    root = schematic.get("root")
    if not root:
        return 1.0
    return compute_node_availability(root, schematic, schematics, overrides, set())


def evaluate_project(json_file, overrides=None):
    with open(json_file, "r") as f:
        data = json.load(f)

    schematics = data["schematics"]
    overrides = overrides or {}

    results = {}
    top_availability = None

    for sid, schematic in schematics.items():
        name = schematic["schematic_name"].strip()
        avail = compute_schematic_availability(schematic, schematics, overrides)
        results[name] = avail
        if name == "- TJB - Unit 3 -":
            top_availability = avail

    return results, top_availability


if __name__ == "__main__":
    overrides = {
        "WTP": 0.0,  # Example override
    }

    results, top_avail = evaluate_project("project_70_rbd.json", overrides=overrides)

    print("\n--- Schematic Availabilities ---")
    for sch_name, avail in results.items():
        print(f"{sch_name:25s}: {avail:.6f}")

    print("\n=== TOTAL SYSTEM AVAILABILITY ===")
    if top_avail is not None:
        print(f"Top schematic (- TJB - Unit 3 -): {top_avail:.6f}")
    else:
        print("Top schematic not found in JSON!")
