import json
import pandas as pd
import numpy as np
from estimators import SingleEstimator, JointEstimator
from system_model import SystemModel
import os

class EvaluateEstimators():
    def __init__(self, measurement_path, single_estimators, joint_estimators, system_model, name=None):
        self._measurement_path = measurement_path
        self._single_estimators = single_estimators
        self._joint_estimators = joint_estimators
        self._system_model = system_model

        self._name = name if name is not None else ""

        self.setup()

    def setup(self):

        i = 0
        while os.path.exists(result_path := f"{self._name}_results_{i}"):
            i += 1
        os.mkdir(result_path)
        self._result_path = result_path

    def run(self):
        with open(self._measurement_path, "r") as f:
            measurements = json.load(f)
            results_benchmark = {}
            for eNodeB, data in measurements.items():
                results_benchmark[eNodeB] = {}

                eNodeB_results_path = os.path.join(self._result_path, eNodeB)
                if not os.path.exists(eNodeB_results_path):
                    os.mkdir(eNodeB_results_path)

                location_list, azimuth_list, pathloss_list = [], [], []

                for sector_index, (sector, sector_data) in enumerate(data["sectors"].items()):
                    results_benchmark[eNodeB][sector] = {}
                    results_benchmark[eNodeB][sector]['metadata'] = {}
                    results_benchmark[eNodeB][sector]['results'] = {}
                    results_benchmark[eNodeB][sector]['metadata']["ground_truth"] = data["ground_truth"][sector_index]

                    locations = np.concatenate(
                        (
                            np.array(sector_data["x"]).reshape(-1, 1),
                            np.array(sector_data["y"]).reshape(-1, 1)
                        ),
                        axis=1
                    )
                    azimuths = np.array(sector_data['azimuth'])
                    pathloss = np.array(sector_data['pathloss'])

                    location_list.append(locations)
                    azimuth_list.append(azimuths)
                    pathloss_list.append(pathloss)

                    self._system_model.plot(
                        locations,
                        pathloss,
                        azimuths,
                        save_path=eNodeB_results_path,
                        name=f"sector_{sector}"
                    )

                    for single_estimator in self._single_estimators:
                        single_estimator.save_path = eNodeB_results_path
                        single_estimator.run(
                            locations,
                            azimuths,
                            pathloss
                        )
                        results_benchmark[eNodeB][sector]['results'][single_estimator.name] = single_estimator.mmse_estimate()
                        single_estimator.traceplot(sector=sector)

                for joint_estimator in self._joint_estimators:
                    joint_estimator.save_path = eNodeB_results_path
                    joint_estimator.run(
                        location_list,
                        azimuth_list,
                        pathloss_list
                    )
                    joint_estimator.traceplot()
                    results = joint_estimator.mmse_estimate()

                    for sector_index, sector in enumerate(data["sectors"].keys()):
                        results_benchmark[eNodeB][sector]['results'][joint_estimator.name] = {}
                        results_benchmark[eNodeB][sector]['results'][joint_estimator.name]["MMSE"] = results[sector_index]
                        results_benchmark[eNodeB][sector]['results'][joint_estimator.name]["sector_decision"] = joint_estimator._w

                with open(os.path.join(self._result_path, "results.json"), "w") as fh_results:
                    json.dump(results_benchmark, fh_results)

if __name__ == "__main__":
    system_model = SystemModel(
        frequency=1.8
    )

    system_model_800 = SystemModel(
        frequency=0.8
    )

    EvaluateEstimators(
        "data/rsrp_pathloss_measurements_1800_1.json",
        [
            SingleEstimator("single_shadow_fading", system_model),
        ],
        [
            JointEstimator("joint_shadow_fading", system_model)
        ],
        system_model,
        "band_1800_1"
    ).run()

    EvaluateEstimators(
        "data/rsrp_pathloss_measurements_1800_2.json",
        [
            SingleEstimator("single_shadow_fading", system_model),
        ],
        [
            JointEstimator("joint_shadow_fading", system_model)
        ],
        system_model,
        "band_1800_2"
    ).run()


    EvaluateEstimators(
        "data/rsrp_pathloss_measurements_800.json",
        [
            SingleEstimator("single_shadow_fading", system_model_800),
        ],
        [
            JointEstimator("joint_shadow_fading", system_model_800)
        ],
        system_model_800,
        "band_800"
    ).run()