# Copyright (c) 2024 Bryan Briney
# Distributed under the terms of the MIT License.
# SPDX-License-Identifier: MIT


import polars as pl


class NoneDict(dict):
    """
    A dictionary that returns None for missing keys.
    """

    def __getitem__(self, key):
        return dict.get(self, key)


schema_dict = {
    "sequence_id": pl.String,
    "v_gene": pl.String,
    "d_gene": pl.String,
    "j_gene": pl.String,
    "c_gene": pl.String,
    "cdr3_length": pl.Int64,
    "junction_aa": pl.String,
    "v_identity": pl.Float64,
    "v_identity_aa": pl.Float64,
    "d_identity": pl.Float64,
    "d_identity_aa": pl.Float64,
    "j_identity": pl.Float64,
    "j_identity_aa": pl.Float64,
    "productive": pl.Boolean,
    "complete_vdj": pl.Boolean,
    "sequence": pl.String,
    "germline": pl.String,
    "sequence_vdjc": pl.String,
    "germline_vdjc": pl.String,
    "sequence_lvdjc": pl.String,
    "germline_lvdjc": pl.String,
    "sequence_aa": pl.String,
    "germline_aa": pl.String,
    "sequence_vdjc_aa": pl.String,
    "germline_vdjc_aa": pl.String,
    "sequence_lvdjc_aa": pl.String,
    "germline_lvdjc_aa": pl.String,
    "sequence_gapped": pl.String,
    "germline_gapped": pl.String,
    # "sequence_gapped_aa": pl.String,  # indels in IMGT gaps can cause translation issues (if part of a codon is on one side of the IMGT gap and the rest is on the other side)
    # "germline_gapped_aa": pl.String,
    "sequence_vdjc_gapped": pl.String,
    "germline_vdjc_gapped": pl.String,
    "sequence_lvdjc_gapped": pl.String,
    "germline_lvdjc_gapped": pl.String,
    "sequence_alignment": pl.String,
    "germline_alignment": pl.String,
    "sequence_alignment_aa": pl.String,
    "germline_alignment_aa": pl.String,
    "umi": pl.String,
    "quality": pl.String,
    "locus": pl.String,
    "species": pl.String,
    "germline_database": pl.String,
    "sequence_input": pl.String,
    "sequence_oriented": pl.String,
    "rev_comp": pl.Boolean,
    "productivity_issues": pl.String,
    "stop_codon": pl.Boolean,
    "v_call": pl.String,
    "v_score": pl.Float64,
    "v_support": pl.Float64,
    "v_cigar": pl.String,
    "v_sequence": pl.String,
    "v_germline": pl.String,
    "v_sequence_gapped": pl.String,
    "v_germline_gapped": pl.String,
    "v_sequence_aa": pl.String,
    "v_germline_aa": pl.String,
    # "v_sequence_gapped_aa": pl.String,
    # "v_germline_gapped_aa": pl.String,
    "v_mutations": pl.String,
    "v_mutations_aa": pl.String,
    "v_mutation_count": pl.Int64,
    "v_mutation_count_aa": pl.Int64,
    "v_insertions": pl.String,
    "v_deletions": pl.String,
    "v_frameshift": pl.Boolean,
    "v_frame": pl.Int64,
    "d_call": pl.String,
    "d_score": pl.Float64,
    "d_support": pl.Float64,
    "d_cigar": pl.String,
    "d_sequence": pl.String,
    "d_germline": pl.String,
    "d_frame": pl.Int64,
    "j_call": pl.String,
    "j_score": pl.Float64,
    "j_support": pl.Float64,
    "j_cigar": pl.String,
    "j_sequence": pl.String,
    "j_germline": pl.String,
    "c_call": pl.String,
    "c_score": pl.Float64,
    "c_support": pl.Float64,
    "c_identity": pl.Float64,
    "c_identity_aa": pl.Float64,
    "c_cigar": pl.String,
    "c_sequence": pl.String,
    "c_germline": pl.String,
    "c_sequence_gapped": pl.String,
    "c_germline_gapped": pl.String,
    "c_sequence_aa": pl.String,
    "c_germline_aa": pl.String,
    "c_sequence_gapped_aa": pl.String,
    "c_germline_gapped_aa": pl.String,
    "c_mutations": pl.String,
    "c_mutations_aa": pl.String,
    "c_mutation_count": pl.Int64,
    "c_mutation_count_aa": pl.Int64,
    "c_insertions": pl.String,
    "c_deletions": pl.String,
    "c_sequence_start": pl.Int64,
    "c_sequence_end": pl.Int64,
    "c_germline_start": pl.Int64,
    "c_germline_end": pl.Int64,
    "c_frame": pl.Int64,
    "np1": pl.String,
    "np2": pl.String,
    "np1_length": pl.Int64,
    "np2_length": pl.Int64,
    "fwr1": pl.String,
    "fwr1_aa": pl.String,
    "cdr1": pl.String,
    "cdr1_aa": pl.String,
    "fwr2": pl.String,
    "fwr2_aa": pl.String,
    "cdr2": pl.String,
    "cdr2_aa": pl.String,
    "fwr3": pl.String,
    "fwr3_aa": pl.String,
    "cdr3": pl.String,
    "cdr3_aa": pl.String,
    "fwr4": pl.String,
    "fwr4_aa": pl.String,
    "junction": pl.String,
    "v_sequence_start": pl.Int64,
    "v_sequence_end": pl.Int64,
    "v_germline_start": pl.Int64,
    "v_germline_end": pl.Int64,
    "j_sequence_start": pl.Int64,
    "j_sequence_end": pl.Int64,
    "j_germline_start": pl.Int64,
    "j_germline_end": pl.Int64,
}


OUTPUT_SCHEMA = NoneDict(schema_dict)
