# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

"""Export a trained model from the full checkpoint (with optimizer etc.) to
a final checkpoint, with only the model itself. The model is always stored as
half float to gain space, and because this has zero impact on the final loss.
When DiffQ was used for training, the model will actually be quantized and bitpacked."""
from argparse import ArgumentParser
from fractions import Fraction
import logging
from pathlib import Path
import sys
import torch

from demucs import train
from demucs.states import serialize_model, save_with_checksum


logger = logging.getLogger(__name__)


def main():
    logging.basicConfig(level=logging.INFO, stream=sys.stderr)

    parser = ArgumentParser("tools.export", description="Export trained models from XP sigs.")
    parser.add_argument('signatures', nargs='*', help='XP signatures.')
    parser.add_argument('-o', '--out', type=Path, default=Path("release_models"),
                        help="Path where to store release models (default release_models)")
    parser.add_argument('-s', '--sign', action='store_true',
                        help='Add sha256 prefix checksum to the filename.')

    args = parser.parse_args()
    args.out.mkdir(exist_ok=True, parents=True)

    for sig in args.signatures:
        xp = train.main.get_xp_from_sig(sig)
        name = train.main.get_name(xp)
        logger.info('Handling %s/%s', sig, name)

        out_path = args.out / (sig + ".th")

        solver = train.get_solver_from_sig(sig)
        if len(solver.history) < solver.args.epochs:
            logger.warning(
                'Model %s has less epoch than expected (%d / %d)',
                sig, len(solver.history), solver.args.epochs)

        solver.model.load_state_dict(solver.best_state)
        pkg = serialize_model(solver.model, solver.args, solver.quantizer, half=True)
        if getattr(solver.model, 'use_train_segment', False):
            batch = solver.augment(next(iter(solver.loaders['train'])))
            pkg['kwargs']['segment'] = Fraction(batch.shape[-1], solver.model.samplerate)
            print("Override", pkg['kwargs']['segment'])
        valid, test = None, None
        for m in solver.history:
            if 'valid' in m:
                valid = m['valid']
            if 'test' in m:
                test = m['test']
        pkg['metrics'] = (valid, test)
        if args.sign:
            save_with_checksum(pkg, out_path)
        else:
            torch.save(pkg, out_path)


if __name__ == '__main__':
    main()