from qgis.core import QgsVectorLayer
import os.path

def get_all_upstream(subby, upstream):
    if subby not in upstream:
        return []

    out = []
    seen = set()
    stack = [(subby, iter(upstream.get(subby, [])))]
    path = [subby]
    on_path = {subby}

    while stack:
        node, it = stack[-1]
        try:
            child = next(it)
        except StopIteration:
            stack.pop()
            on_path.discard(node)
            path.pop()
            continue

        if child in on_path:
            cycle_path = " -> ".join(list(reversed(path + [child])))
            raise ValueError(f"Cycle detected in routing: {cycle_path}")

        if child not in seen:
            seen.add(child)
            out.append(child)
            stack.append((child, iter(upstream.get(child, []))))
            path.append(child)
            on_path.add(child)

    return out

def convert_to_URBS(outputDir,subcatsFile,streamsFile):
    
        
    streamLengths = {}
    streamSlopes = {}
    
    if streamsFile:
        streamsLayer = QgsVectorLayer(streamsFile,"Streams")
        for feature in streamsLayer.getFeatures():
            streamLengths[str(feature["ID"])] = feature["Length_km"]
            streamSlopes[str(feature["ID"])] = feature["Slope_m_m"]
    
    
    subcatsLayer = QgsVectorLayer(subcatsFile,"Subcats")
    required = {"Area_km2", "ID", "Downstream"}
    existing = {f.name() for f in subcatsLayer.fields()}
    missing = sorted(required - existing)
    if missing:
        return "Subcatchments layer is missing required field(s): " + ", ".join(missing)
    subcatAreas = {}
    subcatCatchSlopes = {}
    dsSubcat = {}
    rootedSubbies = {}
    upstream = {}
    originalUpstream = {}
    upstream_counts = {}
    with open(os.path.join(outputDir, "URBS_SubcatFile.csv"), 'w') as f:
        f.write("Index,Area,U,UF,CS,I\n")
        features = sorted(subcatsLayer.getFeatures(), key=lambda f: f["ID"])
        for feature in features:
            area = 0
            fracI = 0
            fracU = 0
            fracUF = 0
            CatSlope = 0
            try:
                index = int(feature['ID'])
            except:
                return "Cannot Determine ID Field for a feature, check feature attributes"
            try:
                area = float(feature['Area_km2'])
            except:
                area = 0
            try:
                fracU = float(feature['FracUrban'])
            except:
                fracU = 0
            try:
                fracUF = float(feature['FracForest'])
            except:
                fracUF = 0
            try:
                fracI = float(feature['FracImp'])
            except:
                fracI = 0
            try:
                CatSlope = float(feature['CatSlope'])
            except:
                CatSlope = 0    
            f.write(f"{feature['ID']},{round(float(area),5)},{round(fracU,5)},{round(fracUF,5)},{round(CatSlope,5)},{round(fracI,5)}\n")
            
            if not streamsFile:
                streamLengths[str(feature["ID"])] = 0.0
                streamSlopes[str(feature["ID"])] = 0.0
                
            dsSubcat[str(feature['ID'])] = str(feature['Downstream'])
            rootedSubbies[str(feature["ID"])] = False
            upstream[str(feature["ID"])] = []
            originalUpstream[str(feature["ID"])] = []
            upstream_counts[str(feature["ID"])] = 0
    
    for subby, downstream in dsSubcat.items():
        if downstream == '-1':
            continue
        try:
            upstream[downstream].append(str(subby))
            originalUpstream[downstream].append(str(subby))
        except:
            return f"Something failed regarding Subarea ID / Stream ID US:{subby} & DS:{downstream}"
    
    usSubbies = {}
    for subby in set(dsSubcat.values()):
        try:
            usSubbies[subby] = get_all_upstream(subby, upstream)
        except Exception as e:
            return str(e)
        upstream_counts[subby] = len(usSubbies[subby])
        
    with open(os.path.join(outputDir, "URBS_RoutingFile.vec"), 'w') as f:
        f.write("MODELNAME\n")
        f.write("Model: SPLIT\n")
        f.write("USES: L , CS , Sc , U , I , F*0.5 \n")
        f.write("DEFAULT PARAMETERS: alpha = 0.005 m = 0.8 beta = 2.5 n = 1 x = 0 IL = 0 CL = 0.0\n")
        f.write("CATCHMENT DATA FILE = URBS_SubcatFile.csv\n")
        upstreamCatch = True
        outlets = 0
        indentCount = 0
        error = False
        while upstreamCatch:
            queue = []
            upstreamCatch = False
            maxSubbiesUs = -99999
            for sub, root_con in rootedSubbies.items():
                if root_con == True:
                    continue
                if upstream_counts[str(sub)] > maxSubbiesUs:
                    queue = [str(sub)]
                    maxSubbiesUs = upstream_counts[str(sub)]
            while len(queue) > 0:
                sub = queue[0]
                if len(upstream[sub]) > 0:
                    maxSubbiesUs = -99999
                    maxSubby = None
                    for usSubby in upstream[sub]:
                        if upstream_counts[str(usSubby)] >= maxSubbiesUs:
                            maxSubby = usSubby
                            maxSubbiesUs = upstream_counts[str(usSubby)]
                    queue.append(maxSubby)
                    queue.pop(0)
                    continue
                
                indentation = ""
                indentation += "\t" * indentCount
                rootedSubbies[queue[0]] = True
                try:
                    streamSlopes[sub]
                except:
                    return f"Stream for subarea {sub}, Does not exist"
                if upstream_counts[sub] == 0:
                    extension = f" L = {round(streamLengths[sub], 5)} Sc = {round(max(streamSlopes[sub], 0.0005), 5)} \n"
                    f.write(f"{indentation}RAIN #{sub}" + extension)
                else:
                    extension = f" L = {round(streamLengths[sub] / 2, 5)} Sc = {round(max(streamSlopes[sub], 0.0005), 5)} \n"
                    f.write(f"{indentation}ADD RAIN #{sub}" + extension)
                
                if (dsSubcat[sub] != '-1' and dsSubcat[sub] != '0'):
                    queue.append(dsSubcat[sub])
                    try:
                        upstream[dsSubcat[sub]].remove(sub)
                    except:
                        return "ERROR DETERMINING UPSTREAM / DOWNSTREAM SUBCATCHMENTS CHECK VECTOR FILE"

                    
                queue.pop(0)
                if len(queue) > 0:
                    if len(upstream[queue[0]]) > 0:
                        f.write(f"{indentation}STORE.\n")
                        indentCount += 1
                    else:
                        for _ in range(1, len(originalUpstream[queue[0]])):
                            indentCount -= 1
                            indentation = ""
                            indentation += "\t" * indentCount
                            f.write(f"{indentation}GET.\n")
                        
                        try:
                            streamSlopes[queue[0]]
                        except:
                            return f"Stream for subarea {queue[0]}, Does not exist"
                        extension = f" L = {round(streamLengths[queue[0]] / 2, 5)} Sc = {round(max(streamSlopes[queue[0]], 0.0005), 5)} \n"
                        f.write(f"{indentation}ROUTE THRU #{queue[0]}" + extension)
            if error:
                break
            for sub, root_con in rootedSubbies.items():
                if root_con == False:
                    f.write("STORE.\n")
                    outlets += 1
                    upstreamCatch = True
                    break
                    
        for _ in range(outlets):
            f.write("GET.\n")
        
        f.write("END OF CATCHMENT DATA.\n")
        return None
    