import random
import json
import sys
from datetime import datetime
from node import Node
from leach import LEACH
from leach_c import LEACHC
from config import FIELD_WIDTH, FIELD_HEIGHT, INITIAL_ENERGY, SCENARIOS, get_num_rounds_for_scenario, ENABLE_MOBILITY
import config


class Simulator:

    def __init__(self, scenario, use_hybrid=False):
        self.scenario = scenario
        self.packet_size = scenario["l"]
        self.probability_ch = scenario["p"]
        self.num_nodes = scenario["n"]
        self.scenario_name = scenario["name"]
        self.use_hybrid = use_hybrid
        self.results = {}
        self.nodes = []

    def initialize_nodes(self):
        self.nodes = []
        for i in range(self.num_nodes):
            x = random.uniform(0, FIELD_WIDTH)
            y = random.uniform(0, FIELD_HEIGHT)
            self.nodes.append(Node(i, x, y, INITIAL_ENERGY))

    def run_protocol(self, protocol_name, protocol_class):
        for node in self.nodes:
            node.energy = INITIAL_ENERGY
            node.is_alive = True
            node.reset_for_round()

        protocol = protocol_class(self.nodes, self.probability_ch, self.packet_size)
        num_rounds = get_num_rounds_for_scenario(self.num_nodes)

        print(f"  Running {protocol_name}...")
        print(f"    Packets: {self.packet_size}, Probability: {self.probability_ch}, Nodes: {self.num_nodes}")

        if self.use_hybrid:
            from simpy_simulator_hybrid import HybridSimPySimulator
            simulator = HybridSimPySimulator(protocol_name, self.nodes, self.packet_size,
                                           self.probability_ch, num_rounds)
            result_data = simulator.run()
            metrics = {
                "final_alive_nodes": sum(1 for n in self.nodes if n.is_alive),
                "first_dead_node_round": result_data["fdn"],
                "first_muted_round": result_data["fmr"],
                "dlbi": result_data["dlbi"],
                "rspi": result_data["rspi"],
                "packets_to_ch": result_data["total_packets_to_ch"],
                "packets_to_bs": result_data["total_packets_to_bs"],
            }
            detailed = result_data.get("rounds_data", [])
        else:
            protocol.run_simulation(num_rounds)
            metrics = protocol.get_metrics(num_rounds)
            detailed = protocol.get_detailed_metrics()

        print(f"  OK - {protocol_name} complete")
        print(f"    Alive: {metrics['final_alive_nodes']}, FDN: {metrics['first_dead_node_round']}")
        print(f"    DLBI: {metrics['dlbi']:.4f}, RSPI: {metrics['rspi']:.4f}")

        return {"protocol": protocol_name, "metrics": metrics, "detailed": detailed}

    def run_simulation(self):
        print(f"\n{'='*60}")
        print(f"Scenario: {self.scenario_name}")
        print(f"{'='*60}")

        self.initialize_nodes()
        self.results["LEACH"] = self.run_protocol("LEACH", LEACH)
        self.results["LEACH-C"] = self.run_protocol("LEACH-C", LEACHC)

        print(f"{'='*60}\n")
        return self.results

    def get_results(self):
        return self.results


def run_all_scenarios(is_static=False, use_hybrid=False):
    config.ENABLE_MOBILITY = not is_static
    all_results = {}
    mode = "STATIC" if is_static else "DYNAMIC"
    approach = "(Hybrid)" if use_hybrid else ""

    print(f"\n{'#'*60}")
    print(f"PHASE: {mode} NETWORKS {approach}")
    print(f"Started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"{'#'*60}\n")

    for scenario in SCENARIOS:
        simulator = Simulator(scenario, use_hybrid=use_hybrid)
        results = simulator.run_simulation()
        all_results[scenario["name"]] = results

    print(f"\n{'#'*60}")
    print(f"Complete: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"{'#'*60}\n")

    return all_results


def save_results(results, output_file):
    json_results = {}
    for scenario_name, scenario_data in results.items():
        json_results[scenario_name] = {
            "LEACH": {
                "metrics": scenario_data["LEACH"]["metrics"],
                "detailed_rounds": scenario_data["LEACH"]["detailed"][:20]
            },
            "LEACH-C": {
                "metrics": scenario_data["LEACH-C"]["metrics"],
                "detailed_rounds": scenario_data["LEACH-C"]["detailed"][:20]
            }
        }

    with open(output_file, 'w') as f:
        json.dump(json_results, f, indent=2)

    print(f"Results saved: {output_file}")


if __name__ == "__main__":
    use_hybrid = "--simpy-hybrid" in sys.argv

    random.seed(42)

    print("\n" + "="*70)
    print("PHASE 1: DYNAMIC SIMULATIONS")
    print("="*70)
    dynamic_results = run_all_scenarios(is_static=False, use_hybrid=use_hybrid)
    save_results(dynamic_results, "results/simulation_results_dynamic.json")

    print("\n" + "="*70)
    print("PHASE 2: STATIC SIMULATIONS")
    print("="*70)
    random.seed(42)
    static_results = run_all_scenarios(is_static=True, use_hybrid=use_hybrid)
    save_results(static_results, "results/simulation_results_static.json")

    print("\n" + "="*70)
    print("COMPARISON: DYNAMIC vs STATIC")
    print("="*70)

    for scenario_name in SCENARIOS[0:1]:
        scenario_label = scenario_name['name']
        print(f"\n{scenario_label}:")
        print("-" * 70)

        for protocol in ["LEACH", "LEACH-C"]:
            dynamic_metrics = dynamic_results[scenario_label][protocol]["metrics"]
            static_metrics = static_results[scenario_label][protocol]["metrics"]

            dyn_fdn = dynamic_metrics['first_dead_node_round'] or "N/A"
            stat_fdn = static_metrics['first_dead_node_round'] or "N/A"
            dyn_fmr = dynamic_metrics['first_muted_round'] or "N/A"
            stat_fmr = static_metrics['first_muted_round'] or "N/A"

            print(f"\n  {protocol}:")
            print(f"    Metric             | Dynamic | Static  | Difference")
            if isinstance(dyn_fdn, int) and isinstance(stat_fdn, int):
                print(f"    FDN                | {dyn_fdn:7} | {stat_fdn:7} | {stat_fdn - dyn_fdn:10.0f}")
            else:
                print(f"    FDN                | {str(dyn_fdn):7} | {str(stat_fdn):7} | N/A")
            if isinstance(dyn_fmr, int) and isinstance(stat_fmr, int):
                print(f"    FMR                | {dyn_fmr:7} | {stat_fmr:7} | {stat_fmr - dyn_fmr:10.0f}")
            else:
                print(f"    FMR                | {str(dyn_fmr):7} | {str(stat_fmr):7} | N/A")
            print(f"    DLBI               | {dynamic_metrics['dlbi']:7.4f} | {static_metrics['dlbi']:7.4f} | {static_metrics['dlbi'] - dynamic_metrics['dlbi']:10.4f}")
            print(f"    RSPI               | {dynamic_metrics['rspi']:7.4f} | {static_metrics['rspi']:7.4f} | {static_metrics['rspi'] - dynamic_metrics['rspi']:10.4f}")

    print(f"\nResults saved to results/ directory")
