from pathlib import Path
from typing import List, Dict
from rdflib import URIRef
from .cim_pydantic.models import (
    Substation, VoltageLevel, ConnectivityNode, TopologicalNode, 
    EnergyConsumer, SynchronousMachine, GeneratingUnit, 
    ACLineSegment, Terminal, BaseVoltage,
    PowerTransformer, PowerTransformerEnd, LinearShuntCompensator,
    Simple_Float
)
from .cim_pydantic.base import serialize_to_turtle

IEEE14_METADATA = {
    "subject": "urn:aceiot-cim-rdf:dataset:ieee-14-bus",
    "properties": {
        "http://purl.org/dc/terms/title": "IEEE 14-Bus System",
        "http://purl.org/dc/terms/source": URIRef(
            "https://electricgrids.engr.tamu.edu/electric-grid-test-cases/ieee-14-bus-system/"
        ),
        "http://purl.org/dc/terms/creator": "Texas A&M University researchers",
        "http://purl.org/dc/terms/description": (
            "The IEEE 14-bus test case represents a simple approximation of the "
            "American Electric Power system as of February 1962. It has 14 buses, "
            "5 generators, and 11 loads."
        ),
        "http://purl.org/dc/terms/rights": (
            "Provided free for commercial or non-commercial use by Texas A&M "
            "University researchers; users are encouraged to cite the appropriate "
            "papers when using the dataset."
        ),
        "http://purl.org/dc/terms/bibliographicCitation": [
            (
                "Electric Grid Test Case Repository, Texas A&M University, "
                "IEEE 14-Bus System."
            ),
            (
                "Power Systems Test Case Archive, University of Washington, "
                "14 Bus Power Flow Test Case."
            ),
        ],
        "http://purl.org/dc/terms/references": URIRef(
            "http://labs.ece.uw.edu/pstca/pf14/pg_tca14bus.htm"
        ),
    },
}

def metadata_for_raw_source(file_path: str):
    name = Path(file_path).name.lower()
    if "ieee 14" in name or "ieee14" in name:
        return IEEE14_METADATA
    return None

def parse_raw(file_path: str):
    with open(file_path, "r") as f:
        lines = [line.strip() for line in f.readlines()]
    
    data = {"buses": [], "loads": [], "shunts": [], "generators": [], "branches": [], "transformers": []}
    line_iter = iter(lines)
    
    # Skip header (3 lines)
    for _ in range(3):
        next(line_iter, None)
    
    def get_parts(line: str) -> List[str]:
        # Split by comma but handle potential spaces/quotes
        return [p.strip().strip("'") for p in line.split(',')]

    # Section 1: Buses
    for line in line_iter:
        if line.startswith("0") and "END OF BUS DATA" in line:
            break
        parts = get_parts(line)
        if parts and parts[0]:
            data["buses"].append({
                "num": parts[0],
                "name": parts[1],
                "base_kv": float(parts[2])
            })

    # Section 2: Loads
    for line in line_iter:
        if line.startswith("0") and "END OF LOAD DATA" in line:
            break
        parts = get_parts(line)
        if parts and parts[0]:
            data["loads"].append({
                "bus": parts[0],
                "p": float(parts[5]),
                "q": float(parts[6])
            })

    # Section 3: Shunts
    for line in line_iter:
        if line.startswith("0") and "END OF FIXED SHUNT DATA" in line:
            break
        parts = get_parts(line)
        if parts and parts[0]:
            data["shunts"].append({
                "bus": parts[0],
                "g": float(parts[3]),
                "b": float(parts[4])
            })

    # Section 4: Generators
    for line in line_iter:
        if line.startswith("0") and "END OF GENERATOR DATA" in line:
            break
        parts = get_parts(line)
        if parts and parts[0]:
            data["generators"].append({
                "bus": parts[0],
                "p": float(parts[2]),
                "q": float(parts[3])
            })

    # Section 5: Branches
    for line in line_iter:
        if line.startswith("0") and "END OF BRANCH DATA" in line:
            break
        parts = get_parts(line)
        if parts and parts[0]:
            data["branches"].append({
                "from_bus": parts[0],
                "to_bus": parts[1],
                "r": float(parts[3]),
                "x": float(parts[4]),
                "b": float(parts[5])
            })

    # Section 6: Transformers
    while True:
        line = next(line_iter, None)
        if not line or (line.startswith("0") and "END OF TRANSFORMER DATA" in line):
            break
        
        parts1 = get_parts(line)
        if not parts1 or not parts1[0]: continue
        
        # 2-winding transformer has 4 lines in v33
        line2 = next(line_iter)
        line3 = next(line_iter)
        line4 = next(line_iter)
        
        parts2 = get_parts(line2)
        parts3 = get_parts(line3)
        parts4 = get_parts(line4)
        
        data["transformers"].append({
            "bus_i": parts1[0],
            "bus_j": parts1[1],
            "bus_k": parts1[2], # 0 for 2-winding
            "r12": float(parts2[0]),
            "x12": float(parts2[1]),
            "sbase12": float(parts2[2]),
            "windv1": float(parts3[0]),
            "nomv1": float(parts3[1]),
            "windv2": float(parts4[0]),
            "nomv2": float(parts4[1])
        })
            
    return data

def convert_to_cim(raw_data: Dict):
    objects = []
    bus_map = {} # bus_num -> (substation, voltage_level, connectivity_node)
    
    # Base Voltages
    base_voltages = {} # kv -> BaseVoltage object
    
    for bus in raw_data["buses"]:
        kv = bus["base_kv"]
        if kv not in base_voltages:
            bv = BaseVoltage(nominalVoltage=kv, mRID=f"bv_{kv}")
            base_voltages[kv] = bv
            objects.append(bv)
            
        sub = Substation(name=f"Substation {bus['num']}", mRID=f"sub_{bus['num']}")
        vl = VoltageLevel(name=f"VL {bus['num']}", mRID=f"vl_{bus['num']}", Substation=sub, BaseVoltage=base_voltages[kv])
        cn = ConnectivityNode(name=f"CN {bus['num']}", mRID=f"cn_{bus['num']}", ConnectivityNodeContainer=vl)
        tn = TopologicalNode(name=f"TN {bus['num']}", mRID=f"tn_{bus['num']}", ConnectivityNodeContainer=vl)
        
        objects.extend([sub, vl, cn, tn])
        bus_map[bus["num"]] = (sub, vl, cn, tn)
        
    for load in raw_data["loads"]:
        _, vl, cn, _ = bus_map[load["bus"]]
        ec = EnergyConsumer(name=f"Load {load['bus']}", mRID=f"load_{load['bus']}", p=load["p"], q=load["q"], EquipmentContainer=vl)
        term = Terminal(name=f"Term Load {load['bus']}", mRID=f"term_load_{load['bus']}", ConductingEquipment=ec, ConnectivityNode=cn)
        objects.extend([ec, term])
        
    for gen in raw_data["generators"]:
        _, vl, cn, _ = bus_map[gen["bus"]]
        gu = GeneratingUnit(name=f"GU {gen['bus']}", mRID=f"gu_{gen['bus']}", EquipmentContainer=vl)
        sm = SynchronousMachine(name=f"Gen {gen['bus']}", mRID=f"gen_{gen['bus']}", p=gen["p"], q=gen["q"], EquipmentContainer=vl, GeneratingUnit=gu)
        term = Terminal(name=f"Term Gen {gen['bus']}", mRID=f"term_gen_{gen['bus']}", ConductingEquipment=sm, ConnectivityNode=cn)
        objects.extend([gu, sm, term])
        
    for i, branch in enumerate(raw_data["branches"]):
        _, vl1, cn1, _ = bus_map[branch["from_bus"]]
        _, vl2, cn2, _ = bus_map[branch["to_bus"]]
        
        line = ACLineSegment(name=f"Line {branch['from_bus']}-{branch['to_bus']}", mRID=f"line_{i}", r=branch["r"], x=branch["x"], bch=branch["b"], EquipmentContainer=vl1)
        term1 = Terminal(name=f"Term 1 Line {i}", mRID=f"term1_line_{i}", ConductingEquipment=line, ConnectivityNode=cn1)
        term2 = Terminal(name=f"Term 2 Line {i}", mRID=f"term2_line_{i}", ConductingEquipment=line, ConnectivityNode=cn2)
        objects.extend([line, term1, term2])

    for i, trans in enumerate(raw_data["transformers"]):
        _, vl1, cn1, _ = bus_map[trans["bus_i"]]
        _, vl2, cn2, _ = bus_map[trans["bus_j"]]
        
        pt = PowerTransformer(name=f"Trans {trans['bus_i']}-{trans['bus_j']}", mRID=f"trans_{i}", EquipmentContainer=vl1)
        
        # PowerTransformerEnd 1 (Side i)
        pte1 = PowerTransformerEnd(name=f"TransEnd 1 {i}", mRID=f"pte1_{i}", PowerTransformer=pt, r=trans["r12"], x=trans["x12"])
        term1 = Terminal(name=f"Term 1 Trans {i}", mRID=f"term1_trans_{i}", ConductingEquipment=pt, ConnectivityNode=cn1)
        pte1.Terminal = term1
        
        # PowerTransformerEnd 2 (Side j)
        pte2 = PowerTransformerEnd(name=f"TransEnd 2 {i}", mRID=f"pte2_{i}", PowerTransformer=pt)
        term2 = Terminal(name=f"Term 2 Trans {i}", mRID=f"term2_trans_{i}", ConductingEquipment=pt, ConnectivityNode=cn2)
        pte2.Terminal = term2
        
        objects.extend([pt, pte1, pte2, term1, term2])
        
    for i, shunt in enumerate(raw_data["shunts"]):
        _, vl, cn, _ = bus_map[shunt["bus"]]
        lsc = LinearShuntCompensator(
            name=f"Shunt {shunt['bus']}", 
            mRID=f"shunt_{i}", 
            gPerSection=shunt["g"], 
            bPerSection=shunt["b"], 
            sections=Simple_Float(value=1.0, mRID=f"shunt_sections_{i}"),
            normalSections=1,
            EquipmentContainer=vl
        )
        term = Terminal(name=f"Term Shunt {shunt['bus']}", mRID=f"term_shunt_{shunt['bus']}", ConductingEquipment=lsc, ConnectivityNode=cn)
        objects.extend([lsc, lsc.sections, term])
        
    return objects

if __name__ == "__main__":
    raw_file = "IEEE 14-Bus System/IEEE 14 bus.raw"
    raw_data = parse_raw(raw_file)
    cim_objects = convert_to_cim(raw_data)
    turtle_data = serialize_to_turtle(cim_objects, metadata=metadata_for_raw_source(raw_file))
    with open("ieee14_cim.ttl", "w") as f:
        f.write(turtle_data)
    print("Exported IEEE 14-bus to ieee14_cim.ttl")
