from aegis_sim.recording.featherrecorder import FeatherRecorder
from aegis_sim.recording.flushrecorder import FlushRecorder
from aegis_sim.recording.phenomaprecorder import PhenomapRecorder
from aegis_sim.recording.picklerecorder import PickleRecorder
from aegis_sim.recording.popgenstatsrecorder import PopgenStatsRecorder
from aegis_sim.recording.progressrecorder import ProgressRecorder
from aegis_sim.recording.summaryrecorder import SummaryRecorder
from aegis_sim.recording.terecorder import TERecorder
from aegis_sim.recording.ticker import Ticker
from aegis_sim.recording.intervalrecorder import IntervalRecorder


def extract_output_specification_from_docstring(method):
    """Extract information about the output file created by the method"""
    docstring = method.__doc__
    texts = docstring.split("# OUTPUT SPECIFICATION")[1:]
    for text in texts:
        parsed = {}
        for pair in text.strip().split("\n"):
            k, v = pair.split(":", maxsplit=1)
            parsed[k.strip()] = v.strip()
        yield parsed


output_specifications_preamble = """
# Output specification

The table below contains information about all files generated by AEGIS. Here we provide an explanation of various table columns.
- `path` gives the location of the stored files. Paths are relative, all files are located in the same folder as the configuration file.
- `filetype` specifies the types of output files. AEGIS uses simple text formats (e.g. .csv, .txt, .json, .yml) and common python filetypes (.pickle and .feather).
- `category` specifies the kind of information stored – information about `genotype`, `phenotype`, `demography`, `population genetics` or `log` (logging information).
- `description` contains short descriptions of the files.
- `trait granularity` specifies the level of trait aggregation: `individual` means that data has not been aggregated across individuals, `popultion` means that the data has been aggregated across individuals, either by computing the count, mean or median of the target trait. Note that when individual-level data is recorded, all individual are included in the recording; not just a sample. In the current version, AEGIS does not keep track of individual IDs, thus individuals are not identifiable across recordings.
- `time granularity` specifies the level of time aggregation: `snapshot` means that no data has been aggregated over multiple simulation steps, while `interval` means that an aggregated metric has been computed across multiple simulation steps. Aggregated information has a smaller storage footprint, allowing it to be stored more frequently. 
- `frequency` specifies the parameter that sets the frequency of recording. Recording of all files is customizable, apart from files that are recorded once (marked by `once`). Higher frequency allows more granular inspection of the simulation but requires more storage space.
- `structure` gives more information about the internal format of the files.

"""

output_specifications = [
    specification
    for method in (
        FeatherRecorder.write_genotypes,
        FeatherRecorder.write_phenotypes,
        FeatherRecorder.write_demography,
        FlushRecorder.write_age_at,
        PhenomapRecorder.write,
        PickleRecorder.write,
        PopgenStatsRecorder.write,
        ProgressRecorder.write_to_progress_log,
        SummaryRecorder.write_input_summary,
        SummaryRecorder.write_output_summary,
        TERecorder.write,
        Ticker.write,
        IntervalRecorder.write_genotypes,
        IntervalRecorder.write_phenotypes,
    )
    for specification in extract_output_specification_from_docstring(
        method=method
    )  # This loop is necessary in case there are multiple output specifications in one method
]
