import argparse
import logging
import os
import sys
import time

import large_image_converter


logger = logging.getLogger('large-image-converter')


def get_parser():
    parser = argparse.ArgumentParser(description='Large Image image converter')
    parser.add_argument(
        '--version', action='version',
        version=large_image_converter.__version__, help='Report version')
    parser.add_argument(
        '--verbose', '-v', action='count', default=0, help='Increase verbosity')
    parser.add_argument(
        '--silent', '-s', action='count', default=0, help='Decrease verbosity')
    parser.add_argument(
        'source', help='Path to source image')
    parser.add_argument(
        'dest', nargs='?', help='Output path')
    parser.add_argument(
        '--compression', '-c',
        choices=[
            '', 'jpeg', 'deflate', 'zip', 'lzw', 'zstd', 'packbits', 'jbig',
            'lzma', 'webp', 'jp2k', 'none',
        ],
        help='Internal compression.  Default will use jpeg if the source '
        'appears to be lossy or lzw if lossless.  lzw is the most compatible '
        'lossless mode.  jpeg is the most compatible lossy mode.  jbig and '
        'lzma may not be available.  jp2k will first write the file with no '
        'compression and then rewrite it with jp2k the specified psnr or '
        'compression ratio.')
    parser.add_argument(
        '--quality', '-q', type=int,
        help='JPEG or webp compression quality.  For webp, specify 0 for '
        'lossless.  Default is 90.')
    parser.add_argument(
        '--level', '-l', type=int,
        help='General compression level.  Used for deflate (zip) (1-9), zstd '
        '(1-22), and some others.')
    parser.add_argument(
        '--predictor', '-p', choices=['', 'none', 'horizontal', 'float', 'yes'],
        help='Predictor for some compressions.  Default is horizontal for '
        'non-geospatial data and yes for geospatial.')
    parser.add_argument(
        '--psnr', type=int,
        help='JP2K peak signal to noise ratio.  0 for lossless.')
    parser.add_argument(
        '--cr', type=int, help='JP2K compression ratio.  1 for lossless.')
    parser.add_argument(
        '--tile', '-t', type=int, help='Tile size.  Default is 256.',
        dest='tileSize')
    parser.add_argument(
        '--overwrite', '-w', action='store_true',
        help='Overwrite an existing output file')
    parser.add_argument(
        '--stats', action='store_true', dest='_stats',
        help='Add conversion stats (time and size) to the ImageDescription of '
        'the output file.  This involves writing the file an extra time; the '
        'stats do not include the extra write.')
    parser.add_argument(
        '--stats-full', '--full-stats',
        action='store_const', const='full', dest='_stats',
        help='Add conversion stats, including noise metrics (PSNR, etc.) to '
        'the output file.  This takes more time and temporary disk space.')
    return parser


def compute_error_metrics(original, altered, results, converterOpts=None):
    """
    Compute the amount of error introduced via conversion compared to
    conversion using a lossless method.  Note that this is not compared to the
    original, as we may not be able to read that in an efficient way.  This is
    a very time-consuming way to compute error metrics, since it first
    reprocesses the input file to a lossless format, then steps through each
    tile and computes RSME and SSIM errors per tile, producing a weighted-by-
    number-of-pixels of each of these.  The RSME is used to compute a PSNR
    value.

    :param original: the original file path.
    :param altered: the path of compressed file to compare.
    :param results: a dictionary to store results.  Modified.
    :param converterOpts: an optional dictionary of parameters used for the
        original conversion.  Only parameters that would affect the selected
        pixels are used.
    """
    import math
    from tempfile import TemporaryDirectory
    import skimage.metrics
    import numpy
    import large_image

    lastlog = 0
    with TemporaryDirectory() as tempDir:
        tempPath = os.path.join(tempDir, os.path.basename(original))
        orig = large_image_converter.convert(original, tempPath, compression='lzw')
        tsOrig = large_image.getTileSource(orig)
        numFrames = len(tsOrig.getMetadata().get('frames', [0]))
        tsAlt = large_image.getTileSource(altered)
        mse = 0
        ssim = 0
        ssim_count = 0
        maxval = 0
        maxdiff = 0
        sum = 0
        count = 0
        tileSize = 2048
        for frame in range(numFrames):
            tiAlt = tsAlt.tileIterator(tile_size=dict(width=tileSize), frame=frame)
            for tileOrig in tsOrig.tileIterator(tile_size=dict(width=tileSize), frame=frame):
                tileAlt = next(tiAlt)
                do = tileOrig['tile'].astype(int)
                da = tileAlt['tile'].astype(int)
                maxval = max(maxval, do.max(), da.max())
                if do.shape[2] > da.shape[2]:
                    do = do[:, :, :da.shape[2]]
                if da.shape[2] > do.shape[2]:
                    da = da[:, :, :do.shape[2]]
                diff = numpy.absolute(do - da)
                maxdiff = max(maxdiff, diff.max())
                sum += diff.sum()
                count += diff.size
                last_mse = numpy.mean(diff ** 2)
                mse += last_mse * diff.size
                last_ssim = 0
                try:
                    last_ssim = skimage.metrics.structural_similarity(
                        do.astype(float), da.astype(float),
                        gaussian_weights=True, sigma=1.5, use_sample_covariance=False,
                        multichannel=do.shape[2] > 1)
                    ssim += last_ssim * diff.size
                    ssim_count += diff.size
                except ValueError:
                    pass
                if time.time() - lastlog >= 10:
                    logger.debug(
                        'Calculating error (%d/%d): rmse %4.2f ssim %6.4f  '
                        'last rmse %4.2f ssim %6.4f' % (
                            tileOrig['tile_position']['position'] + 1,
                            tileOrig['iterator_range']['position'],
                            (mse / count) ** 0.5, ssim / ssim_count,
                            last_mse ** 0.5, last_ssim))
                    lastlog = time.time()
        results['maximum_error'] = maxdiff
        results['average_error'] = sum / count
        results['rmse'] = (mse / count) ** 0.5
        results['psnr'] = 10 * math.log10(
            maxval ** 2 / (mse / count)) if mse else None
        if ssim_count:
            results['ssim'] = ssim / ssim_count
            logger.debug('Calculated error: rmse %4.2f psnr %3.1f ssim %6.4f' % (
                results['rmse'], results['psnr'] or 0, results['ssim']))


def main(args=sys.argv[1:]):
    parser = get_parser()
    opts = parser.parse_args(args=args)
    logger = logging.getLogger('large-image-converter')
    if not len(logger.handlers):
        logger.addHandler(logging.StreamHandler(sys.stderr))
    logger.setLevel(max(1, logging.WARNING - (opts.verbose - opts.silent) * 10))
    try:
        import large_image

        li_logger = large_image.config.getConfig('logger')
        li_logger.setLevel(max(1, logging.CRITICAL - (opts.verbose - opts.silent) * 10))
    except ImportError:
        pass
    logger.debug('Command line options: %r' % opts)
    if not os.path.isfile(opts.source):
        logger.error('Source is not a file (%s)', opts.source)
        return 1
    if opts.compression == 'zip':
        opts.compression = 'deflate'
    converterOpts = {
        k: v for k, v in vars(opts).items()
        if k not in {'source', 'dest', 'verbose', 'silent'} and v is not None}
    start_time = time.time()
    dest = large_image_converter.convert(opts.source, opts.dest, **converterOpts)
    end_time = time.time()
    if not os.path.isfile(dest):
        logger.error('Failed to generate file')
        return 1
    logger.info('Created %s, %d bytes, %3.1f s', dest, os.path.getsize(dest), end_time - start_time)
    if opts._stats:
        import json
        import tifftools.commands

        info = tifftools.read_tiff(dest)
        desc = json.loads(info['ifds'][0]['tags'][tifftools.Tag.ImageDescription.value]['data'])
        desc['large_image_converter']['conversion_stats'] = {
            'time': end_time - start_time,
            'filesize': os.path.getsize(dest),
            'original_filesize': os.path.getsize(opts.source),
            'compression_ratio':
                desc['large_image_converter'].get('frames', 1) *
                sum(info['ifds'][0]['tags'][tifftools.Tag.BitsPerSample.value]['data']) / 8 *
                info['ifds'][0]['tags'][tifftools.Tag.ImageWidth.value]['data'][0] *
                info['ifds'][0]['tags'][tifftools.Tag.ImageLength.value]['data'][0] /
                os.path.getsize(dest),
        }
        if opts._stats == 'full' and opts.compression not in {
                'deflate', 'zip', 'lzw', 'zstd', 'packbits'}:
            compute_error_metrics(
                opts.source, dest, desc['large_image_converter']['conversion_stats'],
                converterOpts)
        tifftools.commands.tiff_set(dest, overwrite=True, setlist=[(
            'ImageDescription', json.dumps(
                desc, separators=(',', ':'), sort_keys=True,
                default=large_image_converter.json_serial))])


if __name__ == '__main__':
    sys.exit(main())
