import json


class JsonWriter:
    def __init__(self, solutions, filename, logger):
        self.solutions = solutions
        self.filename = filename
        self.data = self.prepare_file(solutions)
        self.logger = logger

    def save(self):
        with open(self.filename, 'w') as f:
            json.dump(self.data, f)

    def prepare_file(self, solutions):
        # Initialize an empty dictionary to hold the data
        stations_dict = {}

        # Iterate over the solutions list
        for solution in solutions:
            station_name = solution[0]
            parameters = {"mo": solution[1].mo, "fo": solution[1].fo, "q": solution[1].q, "source_size": solution[1].source_size, "stress_drop": solution[1].stress_drop}
            freq = solution[2]
            amp_th_sp_q = solution[3]
            amp_spectrum = solution[4]

            # Round the parameters to two decimal places
            parameters = {key: round(value, 2) for key, value in parameters.items()}

            # Check if the station already exists in the dictionary
            if station_name not in stations_dict:
                stations_dict[station_name] = {
                    "parameters": list(),
                    "frequency": list(),
                    "fitted_amplitude_spectrum": list(),
                    "amplitude_spectrum": list()
                }

            try:
                # Append the data to the respective lists

                stations_dict[station_name]["parameters"]=parameters
                if freq is not None and amp_spectrum is not None and amp_th_sp_q is not None:
                    stations_dict[station_name]["frequency"].extend(freq)
                    stations_dict[station_name]["fitted_amplitude_spectrum"].extend(amp_th_sp_q)
                    stations_dict[station_name]["amplitude_spectrum"].extend(amp_spectrum)
            except AttributeError as ae:
                self.logger.error(ae)
        return stations_dict