"""Derive a CTC-only encoder by removing the unused TDT projection output.

No weight values or precision settings change. CoreML compilation prunes the
unreferenced projection operations. Probability and length outputs stay intact.
"""
import argparse
import hashlib
import json
from pathlib import Path
import shutil

import coremltools as ct


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    assert not args.out.exists()
    args.out.mkdir(parents=True)
    source = ct.models.MLModel(str(args.source / "Encoder.mlpackage"), skip_model_load=True)
    spec = source.get_spec()
    descriptions = [spec.description, *spec.description.functions]
    removed = []
    for description in descriptions:
        for index in reversed(range(len(description.output))):
            if description.output[index].name == "encoder_proj":
                removed.append(getattr(description, "name", "default"))
                del description.output[index]
    for name, function in spec.mlProgram.functions.items():
        for specialization, block in function.block_specializations.items():
            outputs = [x for x in block.outputs if x != "encoder_proj"]
            assert len(outputs) == 2, (name, specialization, list(block.outputs))
            del block.outputs[:]
            block.outputs.extend(outputs)
    assert removed
    result = ct.models.MLModel(spec, weights_dir=source.weights_dir, skip_model_load=True)
    result.save(str(args.out / "Encoder.mlpackage"))
    shutil.copytree(args.source / "Preprocessor.mlpackage", args.out / "Preprocessor.mlpackage")
    shutil.copyfile(args.source / "vocab.json", args.out / "vocab.json")
    config = json.loads((args.source / "config.json").read_text())
    config.pop("tdt_decoder", None)
    config["models"]["Encoder.mlmodelc"]["outputs"].pop("encoder_proj", None)
    config["models"].pop("DecoderJoint.mlmodelc", None)
    config["models"].pop("JointBlock.mlmodelc", None)
    config["model"] = "parakeet-ctc-110m-coreml"
    config["version"] = "1.1.0-ctc-candidate"
    (args.out / "config.json").write_text(json.dumps(config, indent=2) + "\n")
    before = Path(source.weights_dir) / "weight.bin"
    after = next((args.out / "Encoder.mlpackage").rglob("weight.bin"))
    assert hashlib.sha256(before.read_bytes()).digest() == hashlib.sha256(after.read_bytes()).digest()
    (args.out / "derivation.json").write_text(json.dumps({"removed_output": "encoder_proj", "functions": sorted(spec.mlProgram.functions),
        "encoder_weight_sha256": hashlib.sha256(after.read_bytes()).hexdigest(), "weight_bytes_unchanged": True}, indent=2) + "\n")
    print("CTC-only encoder derived with unchanged weights:", args.out)


if __name__ == "__main__":
    main()