import numpy as np
import pandas as pd
from scipy.spatial import distance_matrix
from scipy.sparse.csgraph import minimum_spanning_tree
from collections import defaultdict
import random

try:
    import pulp

    pulp_available = True
except ImportError:
    pulp = None
    pulp_available = False

try:
    import pyomo.environ as pyo_env

    pyomo_available = True
except (ImportError, AttributeError):
    pyomo_available = False
    print("Pyomo not available, falling back to PuLP")


def design_with_pyomo(
    turbines,
    substation,
    cable_specs=None,
    voltage=66000,
    power_factor=0.95,
    system_params=None,
    max_clusters=None,
    time_limit=300,
    evaluate_func=None,
    total_invest_func=None,
    get_max_capacity_func=None,
):
    """
    使用Pyomo求解器优化集电线路布局
    :param turbines: 风机DataFrame
    :param substation: 升压站坐标
    :param cable_specs: 电缆规格
    :param system_params: 系统参数（用于NPV计算）
    :param max_clusters: 最大簇数，默认基于功率计算
    :param time_limit: 求解时间限制（秒）
    :param evaluate_func: 评估函数
    :param total_invest_func: 总投资计算函数
    :param get_max_capacity_func: 获取最大容量函数
    :return: 连接列表和带有簇信息的turbines
    """
    if get_max_capacity_func:
        max_mw = get_max_capacity_func(cable_specs, voltage, power_factor)
    else:
        max_mw = 100.0
    total_power = turbines["power"].sum()
    if max_clusters is None:
        max_clusters = int(np.ceil(total_power / max_mw))
    n_turbines = len(turbines)

    all_coords = np.vstack([substation, turbines[["x", "y"]].values])
    dist_matrix_full = distance_matrix(all_coords, all_coords)

    # Simple fallback for now - use PuLP instead
    print("Pyomo not fully implemented, falling back to PuLP")
    return design_with_mip(
        turbines,
        substation,
        cable_specs,
        voltage,
        power_factor,
        system_params,
        max_clusters,
        time_limit,
        evaluate_func,
        total_invest_func,
        get_max_capacity_func,
    )


def design_with_mip(
    turbines,
    substation,
    cable_specs=None,
    voltage=66000,
    power_factor=0.95,
    system_params=None,
    max_clusters=None,
    time_limit=300,
    evaluate_func=None,
    total_invest_func=None,
    get_max_capacity_func=None,
):
    """
    使用混合整数规划(MIP)优化集电线路布局
    :param turbines: 风机DataFrame
    :param substation: 升压站坐标
    :param cable_specs: 电缆规格
    :param system_params: 系统参数（用于NPV计算）
    :param max_clusters: 最大簇数，默认基于功率计算
    :param time_limit: 求解时间限制（秒）
    :param evaluate_func: 评估函数
    :param total_invest_func: 总投资计算函数
    :param get_max_capacity_func: 获取最大容量函数
    :return: 连接列表和带有簇信息的turbines
    """
    if not pulp_available:
        print(
            "WARNING: PuLP library not available. MIP optimization skipped, falling back to MST."
        )
        from main import design_with_mst

        connections = design_with_mst(turbines, substation)
        return connections, turbines

    if get_max_capacity_func:
        max_mw = get_max_capacity_func(cable_specs, voltage, power_factor)
    else:
        max_mw = 100.0
    if max_clusters is None:
        max_clusters = int(np.ceil(turbines["power"].sum() / max_mw))
    n_turbines = len(turbines)

    print(
        f"MIP Model Setup: n_turbines={n_turbines}, max_clusters={max_clusters}, max_mw={max_mw:.2f} MW"
    )

    all_coords = np.vstack([substation, turbines[["x", "y"]].values])
    dist_matrix_full = distance_matrix(all_coords, all_coords)

    prob = pulp.LpProblem("WindFarmCollectorMIP", pulp.LpMinimize)

    # Create all decision variables upfront to avoid duplicates
    assign_vars = {}
    for i in range(n_turbines):
        for k in range(max_clusters):
            assign_vars[(i, k)] = pulp.LpVariable(f"assign_{i}_{k}", cat="Binary")

    cluster_vars = {}
    for k in range(max_clusters):
        cluster_vars[k] = pulp.LpVariable(f"cluster_{k}", cat="Binary")

    # Helper functions to access variables
    def assign_var(i, k):
        return assign_vars[(i, k)]

    def cluster_var(k):
        return cluster_vars[k]

    # Simplified objective function: minimize total distance
    prob += pulp.lpSum(
        [
            dist_matrix_full[0, i + 1] * assign_var(i, k)
            for i in range(n_turbines)
            for k in range(max_clusters)
        ]
    )

    for i in range(n_turbines):
        prob += pulp.lpSum([assign_var(i, k) for k in range(max_clusters)]) == 1

    for k in range(max_clusters):
        cluster_power = pulp.lpSum(
            [turbines.iloc[i]["power"] * assign_var(i, k) for i in range(n_turbines)]
        )
        prob += cluster_power <= max_mw * 1.2 * cluster_var(k)

    for k in range(max_clusters):
        for i in range(n_turbines):
            prob += assign_var(i, k) <= cluster_var(k)

    print(
        f"MIP Model: {len(prob.variables())} variables, {len(prob.constraints)} constraints"
    )

    # Debug: Print model structure
    print("MIP model structure check:")
    print(f"  Variables: {len(prob.variables())}")
    print(f"  Constraints: {len(prob.constraints)}")
    print(f"  Time limit: {time_limit}s")
    print(f"  Turbines: {n_turbines}, Clusters: {max_clusters}")

    # Test solver availability
    try:
        import subprocess

        test_solver = subprocess.run(
            [
                r"D:\code\windfarm\.venv\Lib\site-packages\pulp\apis\..\solverdir\cbc\win\i64\cbc.exe",
                "-version",
            ],
            capture_output=True,
            text=True,
            timeout=5,
        )
        print(
            f"CBC solver test: {test_solver.stdout[:100] if test_solver.stdout else 'No output'}"
        )
    except Exception as solver_test_error:
        print(f"CBC solver test failed: {solver_test_error}")

    print("MIP: Starting to solve...")
    try:
        # Try to use CBC solver with different configurations
        solver = pulp.PULP_CBC_CMD(
            timeLimit=time_limit,
            msg=False,
            warmStart=False,
        )
        print(f"Using CBC solver with time limit: {time_limit}s")
        status = prob.solve(solver)
        print(
            f"MIP: Solver status={pulp.LpStatus[prob.status]}, Objective value={pulp.value(prob.objective):.4f}"
        )
    except Exception as e:
        print(f"MIP: CBC solver execution failed: {e}")
        # Try alternative solver configurations
        try:
            print("MIP: Trying alternative solver configuration...")
            solver = pulp.PULP_CBC_CMD(
                msg=True,  # Enable messages for debugging
                threads=1,  # Single thread
                timeLimit=time_limit,
            )
            status = prob.solve(solver)
            print(
                f"MIP: Alternative solver status={pulp.LpStatus[prob.status]}, Objective value={pulp.value(prob.objective):.4f}"
            )
        except Exception as e2:
            print(f"MIP: All solver attempts failed: {e2}, falling back to MST")
            from main import design_with_mst

            connections = design_with_mst(turbines, substation)
            return connections, turbines

    if pulp.LpStatus[prob.status] != "Optimal":
        print(
            f"MIP solver status: {pulp.LpStatus[prob.status]}, solution not found, falling back to MST"
        )
        print("Model feasibility check:")
        print(f"Total power: {turbines['power'].sum():.2f} MW")
        print(f"Max cluster capacity: {max_mw:.2f} MW")
        print(f"Number of clusters: {max_clusters}, Number of turbines: {n_turbines}")

        for k in range(max_clusters):
            cluster_power = pulp.value(
                pulp.lpSum(
                    [
                        turbines.iloc[i]["power"] * assign_var(i, k)
                        for i in range(n_turbines)
                    ]
                )
            )
            cluster_used = pulp.value(cluster_var(k))
            print(
                f"Cluster {k}: Power={cluster_power:.2f} MW (max {max_mw * 1.2:.2f}), Used={cluster_used}"
            )

        from main import design_with_mst

        connections = design_with_mst(turbines, substation)
        return connections, turbines

    cluster_assign = [-1] * n_turbines
    active_clusters = []
    for k in range(max_clusters):
        if pulp.value(cluster_var(k)) > 0.5:
            active_clusters.append(k)

    for i in range(n_turbines):
        assigned = False
        for k in active_clusters:
            if pulp.value(assign_var(i, k)) > 0.5:
                cluster_assign[i] = k
                assigned = True
                break
        if not assigned:
            dists = [dist_matrix_full[0, i + 1] for k in active_clusters]
            cluster_assign[i] = active_clusters[np.argmin(dists)]

    clusters = defaultdict(list)
    for i, c in enumerate(cluster_assign):
        clusters[c].append(i)

    connections = []
    for c, members in clusters.items():
        if len(members) == 0:
            continue
        coords = turbines.iloc[members][["x", "y"]].values
        if len(members) > 1:
            dm = distance_matrix(coords, coords)
            mst = minimum_spanning_tree(dm).toarray()
            for i in range(len(members)):
                for j in range(len(members)):
                    if mst[i, j] > 0:
                        connections.append(
                            (
                                f"turbine_{members[i]}",
                                f"turbine_{members[j]}",
                                mst[i, j],
                            )
                        )
        dists = [dist_matrix_full[0, m + 1] for m in members]
        closest = members[np.argmin(dists)]
        connections.append((f"turbine_{closest}", "substation", min(dists)))

    turbines["cluster"] = cluster_assign

    # Check cluster distances
    min_cluster_distance = check_cluster_distances(clusters, turbines)
    if min_cluster_distance is not None:
        print(
            f"Cluster validation: Minimum distance between clusters = {min_cluster_distance:.2f} m"
        )
        if min_cluster_distance < 1000:
            print(
                f"WARNING: Clusters are very close to each other ({min_cluster_distance:.2f} m < 1000 m)"
            )
        elif min_cluster_distance < 2000:
            print(
                f"NOTICE: Clusters are relatively close ({min_cluster_distance:.2f} m)"
            )

    # Check for cable crossings
    cable_crossings = check_cable_crossings(connections, turbines, substation)
    if cable_crossings:
        print(
            f"WARNING: Found {len(cable_crossings)} cable crossing(s) in the solution"
        )
        for i, (idx1, idx2, p1, p2, p3, p4) in enumerate(cable_crossings):
            conn1 = connections[idx1]
            conn2 = connections[idx2]
            print(
                f"  Crossing {i + 1}: Connection {conn1[0]}-{conn1[1]} crosses {conn2[0]}-{conn2[1]}"
            )
    else:
        print("No cable crossings detected in the solution")

    print(
        f"MIP optimization completed successfully, {len(connections)} connections generated"
    )
    return connections, turbines


def calculate_cluster_centroids(clusters, turbines):
    """Calculate the centroid coordinates for each cluster."""
    centroids = {}
    for c, members in clusters.items():
        if len(members) == 0:
            centroids[c] = (0, 0)
        else:
            coords = turbines.iloc[members][["x", "y"]].values
            centroid_x = np.mean(coords[:, 0])
            centroid_y = np.mean(coords[:, 1])
            centroids[c] = (centroid_x, centroid_y)
    return centroids


def check_cluster_distances(clusters, turbines, min_distance_threshold=1000):
    """Check if any clusters are too close to each other."""
    if len(clusters) < 2:
        return None

    centroids = calculate_cluster_centroids(clusters, turbines)
    active_clusters = [c for c, members in clusters.items() if len(members) > 0]

    min_distance = float("inf")
    min_pair = None

    for i in range(len(active_clusters)):
        for j in range(i + 1, len(active_clusters)):
            c1, c2 = active_clusters[i], active_clusters[j]
            centroid1 = np.array(centroids[c1])
            centroid2 = np.array(centroids[c2])
            distance = np.linalg.norm(centroid1 - centroid2)

            if distance < min_distance:
                min_distance = distance
                min_pair = (c1, c2)

    return min_distance


def check_cable_crossings(connections, turbines, substation):
    """Check if there are cable crossings in the solution."""
    crossings = []

    def line_intersection(p1, p2, p3, p4):
        """Check if line segments (p1,p2) and (p3,p4) intersect."""
        x1, y1 = p1
        x2, y2 = p2
        x3, y3 = p3
        x4, y4 = p4

        denom = (y4 - y3) * (x2 - x1) - (x4 - x3) * (y2 - y1)

        if abs(denom) < 1e-10:
            return False

        ua = ((x4 - x3) * (y1 - y3) - (y4 - y3) * (x1 - x3)) / denom
        ub = ((x2 - x1) * (y1 - y3) - (y2 - y1) * (x1 - x3)) / denom

        return 0 <= ua <= 1 and 0 <= ub <= 1

    def get_turbine_coord(connection_part):
        """Get coordinates from connection part (turbine_# or substation)."""
        if connection_part == "substation":
            # Handle different substation formats robustly
            if isinstance(substation, np.ndarray):
                if substation.ndim == 1:
                    # 1D array [x, y]
                    return (substation[0], substation[1])
                elif substation.ndim == 2:
                    # 2D array [[x, y]] or shape (n, 2)
                    if substation.shape[0] == 1:
                        return (substation[0, 0], substation[0, 1])
                    else:
                        # Multiple points, use first one
                        return (substation[0, 0], substation[0, 1])
                else:
                    # Unexpected dimension, try fallback
                    return (substation.flat[0], substation.flat[1])
            elif isinstance(substation, (list, tuple)):
                # List or tuple format
                # Handle nested lists like [[x, y]]
                if (
                    isinstance(substation[0], (list, tuple, np.ndarray))
                    and len(substation[0]) >= 2
                ):
                    return (substation[0][0], substation[0][1])
                elif len(substation) >= 2:
                    return (substation[0], substation[1])
                else:
                    return (float("inf"), float("inf"))
            else:
                # Unexpected format, try to convert
                try:
                    sub_array = np.array(substation)
                    if sub_array.ndim == 1:
                        return (sub_array[0], sub_array[1])
                    else:
                        return (sub_array.flat[0], sub_array.flat[1])
                except:
                    return (float("inf"), float("inf"))
        else:
            turbine_idx = int(connection_part.split("_")[1])
            return (
                turbines.iloc[turbine_idx]["x"],
                turbines.iloc[turbine_idx]["y"],
            )

    for i in range(len(connections)):
        for j in range(i + 1, len(connections)):
            conn1 = connections[i]
            conn2 = connections[j]

            p1 = get_turbine_coord(conn1[0])
            p2 = get_turbine_coord(conn1[1])
            p3 = get_turbine_coord(conn2[0])
            p4 = get_turbine_coord(conn2[1])

            if (
                np.array_equal(p1, p3)
                or np.array_equal(p1, p4)
                or np.array_equal(p2, p3)
                or np.array_equal(p2, p4)
            ):
                continue

            if line_intersection(p1, p2, p3, p4):
                crossings.append((i, j, p1, p2, p3, p4))

    return crossings