import pandas as pd
from mdocfile.data_models import Mdoc

from ._threaded import run_threaded


def _tilt_mdoc(
    mdoc_file, tlt_file, skipped_tilts, verbose=False, dry_run=False, overwrite=False
):
    output = mdoc_file.parent / "mdoc_tilted" / mdoc_file.name

    if not overwrite and output.exists():
        raise FileExistsError(output)

    if verbose:
        print(f"Tilting mdoc: {mdoc_file}")
        print(f"using: {tlt_file}")
        print(f"saving to {output}")

    if not dry_run:
        mdoc = Mdoc.from_file(mdoc_file)
        new_angles = iter(pd.read_csv(tlt_file, header=None)[0])
        # mdoc is still out of order, but aretomo puts the new_angles in order
        # this gives us tuples like (original_index, section), but ordered by angle
        # like aretomo so we can compare it with the new angles
        sorted_section_data = sorted(
            enumerate(mdoc.section_data), key=lambda x: x[1].TiltAngle
        )
        for idx, section in sorted_section_data:
            if idx in skipped_tilts:
                continue
            try:
                section.TiltAngle = next(new_angles)
            except StopIteration as e:
                raise RuntimeError(
                    f"not enough tilts generated by aretomo in {tlt_file}, "
                    f"or Mdoc is corrupted"
                ) from e

        try:
            next(new_angles)
        except StopIteration:
            pass
        else:
            raise RuntimeError(f"too many tilts generated by aretomo in {tlt_file}")

        with open(output, "w+" if overwrite else "w") as f:
            f.write(mdoc.to_string())


def tilt_mdocs_batch(progress, tilt_series, **kwargs):
    partials = []
    for ts in tilt_series:
        partials.append(
            lambda ts=ts: _tilt_mdoc(
                mdoc_file=ts["mdoc"],
                tlt_file=ts["tlt"],
                skipped_tilts=ts["skipped_tilts"],
                **kwargs,
            )
        )

    run_threaded(progress, partials, label="Creating tilted mdocs...", **kwargs)
