import argparse
import multiprocessing
import sys
import typing

import numpy as np

from ..imaris import parse_terastitcher
from ..stitch import StitchSrcVolume, do_stitch, adjust_alignments


def first_prime_at_or_before(n:int):
    """Find the first prime at or before the given #"""
    assert n >= 2, f"There are no primes less than {n}"
    for guess in range(n, 1, -1):
        for factor in range(2, int(np.sqrt(guess))+1):
            if guess % factor == 0:
                break
        else:
            return guess
    raise AssertionError(f"There should be a prime less than or equal to {n}")

def parse_arguments(args:typing.Sequence[str]):
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--terastitcher-xml",
        help="XML file generated by microscope for use by Terastitcher",
        required=True
    )
    parser.add_argument(
        "--alignment",
        help="JSON alignment file created by oblique-align",
        required=True
    )
    parser.add_argument(
        "--output",
        help="Output directory for stitched volume",
        required=True
    )
    parser.add_argument(
        "--n-cores",
        help="# of processes to use for input and stitching",
        type=int,
        default=min(multiprocessing.cpu_count(), 24)
    )
    parser.add_argument(
        "--n-writers",
        help="# of processes to use for output (hint - prime #s are better).",
        type=int,
        default=first_prime_at_or_before(
            max(3, min(multiprocessing.cpu_count(), 17)))
    )
    parser.add_argument(
        "--channel",
        help="Channel number to stitch (1 to # of channels in ims file)",
        default=1,
        type=int
    )
    parser.add_argument(
        "--levels",
        help="# of levels in the downsampling pyramid",
        default=5,
        type=int
    )
    parser.add_argument(
        "--output-size",
        help="Size of the output volume (x,y,z). Defaults to the extent of all "
             "prestitched volumes.")
    parser.add_argument(
        "--output-offset",
        help="Offset of the output volume. Only use with --output-size. ")
    return parser.parse_args(args)


VOLUMES = {}


def main(args=sys.argv[1:]):
    opts = parse_arguments(args)
    all_volumes = []
    for k, v in \
             parse_terastitcher(opts.terastitcher_xml).items():
        VOLUMES[k] = v
        all_volumes.append(v)
        v.directory.current_channel = opts.channel - 1
    adjust_alignments(opts, all_volumes)
    StitchSrcVolume.rebase_all(all_volumes, True)
    do_stitch(opts.output,
              all_volumes,
              levels=opts.levels,
              n_workers=opts.n_cores,
              n_writers=opts.n_writers,
              output_size=opts.output_size,
              output_offset=opts.output_offset)


if __name__ == "__main__":
    main()

