"""Gene and transcript representations."""

from dataclasses import dataclass, replace
from functools import cached_property
from typing import Optional, List, Dict, Set, Tuple, Union, Callable
import numpy as np
from weakref import WeakKeyDictionary
from pathlib import Path

from .genes_auxiliary import build_annotated_regions
from .strand import Strand
from .intervals import GenomicInterval, GenomicIntervalsCollection, AnnotatedIntervalSet
from . import genes_auxiliary

_region_cache = WeakKeyDictionary()


@dataclass(frozen=True)
class Transcript:
    blocks: np.ndarray  # (N, 2) array: start, end
    strand: Strand
    chrom: str
    id: str
    cds_start: Optional[int] = None
    cds_end: Optional[int] = None
    biotype: Optional[str] = None

    @classmethod
    def from_intervals_collection(
        cls,
        collection: GenomicIntervalsCollection,
        transcript_id: str,
        cds_start: Optional[int] = None,
        cds_end: Optional[int] = None,
        biotype: Optional[str] = None,
    ) -> 'Transcript':
        """Create a Transcript from a GenomicIntervalsCollection.

        The collection already guarantees same chrom, same strand, and
        sorted-by-start blocks. This method additionally validates that
        blocks do not overlap.
        """
        if collection.is_empty():
            raise ValueError("Cannot create Transcript from empty collection")

        blocks = collection.array
        if len(blocks) > 1:
            ends = blocks[:-1, 1]
            next_starts = blocks[1:, 0]
            if np.any(ends > next_starts):
                first = int(np.argmax(ends > next_starts))
                raise ValueError(
                    f"Overlapping blocks: block {first} ends at {blocks[first, 1]} "
                    f"but block {first + 1} starts at {blocks[first + 1, 0]}"
                )

        return cls(
            blocks=blocks.copy(),
            strand=collection.strand,
            chrom=collection.chrom,
            id=transcript_id,
            cds_start=cds_start,
            cds_end=cds_end,
            biotype=biotype,
        )

    def with_id(self, new_id: str) -> 'Transcript':
        """Return a copy of this transcript with a different ID."""
        return replace(self, id=new_id)

    def with_fields(self, **kwargs) -> 'Transcript':
        """Return a copy with arbitrary fields replaced.

        Example:
            t.with_fields(id="new_id", biotype="lncRNA")
        """
        return replace(self, **kwargs)

    @property
    def is_coding(self) -> bool:
        return self.cds_start is not None and self.cds_end is not None

    def __hash__(self) -> int:
        return hash(self.id)

    @cached_property
    def transcript_span(self) -> np.ndarray:
        """Get genomic span of transcript (min start, max end) regardless of block order."""
        if len(self.blocks) == 0:
            raise ValueError("Cannot compute transcript span for transcript with no blocks")
        return np.array([self.blocks[:, 0].min(), self.blocks[:, 1].max()], dtype=np.int32)

    @cached_property
    def start(self) -> int:
        return int(self.transcript_span[0])

    @cached_property
    def end(self) -> int:
        return int(self.transcript_span[1])

    @cached_property
    def transcript_interval(self) -> GenomicInterval:
        return genes_auxiliary.get_transcript_interval(self)

    @cached_property
    def transcript_cds_interval(self) -> GenomicInterval:
        return genes_auxiliary.get_transcript_cds_interval(self)

    def exons(self) -> np.ndarray:
        return self.blocks

    @cached_property
    def cds_blocks(self) -> np.ndarray:
        return genes_auxiliary.get_cds_blocks(self)

    @cached_property
    def utr5_blocks(self) -> np.ndarray:
        return genes_auxiliary.get_utr5_blocks(self)

    @cached_property
    def utr3_blocks(self) -> np.ndarray:
        return genes_auxiliary.get_utr3_blocks(self)
    
    @cached_property
    def utr_blocks(self) -> np.ndarray:
        """Get all UTR intervals (merged 5' and 3')."""
        utr5 = self.utr5_blocks
        utr3 = self.utr3_blocks
        
        if utr5.size == 0 and utr3.size == 0:
            return np.empty((0, 2), dtype=np.uint32)
        elif utr5.size == 0:
            return utr3
        elif utr3.size == 0:
            return utr5
        else:
            return np.vstack([utr5, utr3])

    def __str__(self) -> str:
        from ..ops.transcript_serialization import transcript_to_bed12_string
        return transcript_to_bed12_string(self)
    
    def __repr__(self) -> str:
        coding_info = f", coding={self.is_coding}" if self.is_coding else ", non-coding"
        cds_info = f", CDS={self.cds_start}-{self.cds_end}" if self.is_coding else ""
        
        strand_str = Strand(self.strand).to_char()
            
        return f"Transcript(id='{self.id}', {self.chrom}:{self.blocks[0,0]}-{self.blocks[-1,1]}:{strand_str}, {len(self.blocks)} exons{coding_info}{cds_info})"

    def compute_flanks(
            self,
            flank_size: int,
            chrom_sizes: Dict[str, int]
    ) -> Tuple[Optional[GenomicInterval], Optional[GenomicInterval]]:
        return genes_auxiliary.compute_flanks(self, flank_size, chrom_sizes)

    def get_annotated_regions(self, chrom_sizes: dict, flank_size: int = 5000) -> AnnotatedIntervalSet:
        cache = _region_cache.setdefault(self, {})
        key = (flank_size, id(chrom_sizes))
        if key not in cache:
            cache[key] = build_annotated_regions(self, chrom_sizes, flank_size)
        return cache[key]
    
    def contains_interval(self, interval: GenomicInterval) -> bool:
        if self.chrom != interval.chrom:
            return False
        
        block_starts = self.blocks[:, 0]
        block_ends = self.blocks[:, 1]
        
        no_overlap = (interval.end <= block_starts) | (block_ends <= interval.start)
        has_overlap = ~no_overlap
        
        return np.any(has_overlap)
    
    def splice_junctions(self):
        """Generator yielding splice junction coordinates (donor, acceptor) for transcript."""
        if len(self.blocks) <= 1:
            # Single exon transcript has no splice junctions
            return
        
        sorted_blocks = self.blocks[np.argsort(self.blocks[:, 0])]
        for i in range(len(sorted_blocks) - 1):
            current_end = sorted_blocks[i, 1]      # Donor site (end of current exon)
            next_start = sorted_blocks[i + 1, 0]   # Acceptor site (start of next exon)
            
            if self.strand == Strand.MINUS:
                yield (next_start, current_end)  # (acceptor, donor)
            else:
                yield (current_end, next_start)  # (donor, acceptor)
                
    def get_introns(self, use_numba: bool = True) -> np.ndarray:
        if len(self.blocks) <= 1:
            return np.empty((0, 2), dtype=np.int32)
        
        intron_starts = self.blocks[:-1, 1]
        intron_ends = self.blocks[1:, 0]
        valid_mask = intron_starts < intron_ends
        
        if not np.any(valid_mask):
            return np.empty((0, 2), dtype=np.int32)
            
        return np.column_stack((intron_starts[valid_mask], intron_ends[valid_mask]))


class Gene:
    """Gene containing multiple transcripts with computed genomic bounds."""
    def __init__(self, gene_id: str, transcripts: List[Transcript], gene_name: Optional[str] = None):
        self.gene_id = gene_id
        self.gene_name = gene_name
        self._transcripts = list(transcripts)
        self._canonical_transcript_id: Optional[str] = None
        
        if not self._transcripts:
            raise ValueError(f"Gene {gene_id} must have at least one transcript")
        
        # Validate all transcripts are on same chromosome and strand
        first_transcript = self._transcripts[0]
        self._chrom = first_transcript.chrom
        self._strand = first_transcript.strand
        
        for transcript in self._transcripts[1:]:
            if transcript.chrom != self._chrom:
                raise ValueError(f"All transcripts must be on same chromosome. Found {transcript.chrom} and {self._chrom}")
            if transcript.strand != self._strand:
                raise ValueError(f"All transcripts must have same strand. Found {transcript.strand} and {self._strand}")
    
    @property
    def transcripts(self) -> List[Transcript]:
        return self._transcripts.copy()
    
    @property
    def transcript_ids(self) -> Set[str]:
        return {t.id for t in self._transcripts}
    
    @property
    def chrom(self) -> str:
        return self._chrom

    @property
    def strand(self) -> Strand:
        return self._strand
    
    @cached_property
    def start(self) -> int:
        return int(min(t.transcript_span[0] for t in self._transcripts)) if self._transcripts else 0
    
    @cached_property
    def end(self) -> int:
        return int(max(t.transcript_span[1] for t in self._transcripts)) if self._transcripts else 0

    @cached_property
    def bounds(self) -> GenomicInterval:
        return GenomicInterval(self.chrom, self.start, self.end, self.strand)

    @property
    def length(self) -> int:
        return self.end - self.start
    
    def get_transcript(self, transcript_id: str) -> Optional[Transcript]:
        for transcript in self._transcripts:
            if transcript.id == transcript_id:
                return transcript
        return None
    
    def has_transcript(self, transcript_id: str) -> bool:
        return transcript_id in self.transcript_ids
    
    @property
    def is_coding(self) -> bool:
        """Check if gene has any coding transcripts."""
        return any(t.is_coding for t in self._transcripts)

    def to_union_transcript(
        self,
        transcript_id: Optional[str] = None,
        id_prefix: str = "U_",
    ) -> Transcript:
        """Merge all isoforms into a single union transcript.

        For coding genes: merges CDS and UTR blocks separately, then combines.
        CDS boundaries are derived from the union of all CDS blocks.
        For non-coding genes: merges all exon blocks.

        Args:
            transcript_id: Explicit ID for the union transcript.
                If None, uses ``id_prefix + gene_id``.
            id_prefix: Prefix prepended to gene_id when transcript_id is None.
        """
        from ..ops.interval_ops import merge_intervals

        tid = transcript_id if transcript_id is not None else f"{id_prefix}{self.gene_id}"
        ts = self._transcripts

        has_coding = any(t.is_coding for t in ts)
        if has_coding:
            cds_arrays = [t.cds_blocks for t in ts if t.is_coding and t.cds_blocks.size > 0]
            utr_arrays = [t.utr_blocks for t in ts if t.is_coding and t.utr_blocks.size > 0]

            cds_union = merge_intervals(np.vstack(cds_arrays)) if cds_arrays else np.empty((0, 2), dtype=np.int32)
            utr_union = merge_intervals(np.vstack(utr_arrays)) if utr_arrays else np.empty((0, 2), dtype=np.int32)

            if cds_union.size == 0 and utr_union.size == 0:
                all_blocks = [t.blocks for t in ts if t.blocks.size > 0]
                merged = merge_intervals(np.vstack(all_blocks)) if all_blocks else np.empty((0, 2), dtype=np.int32)
                cds_start, cds_end = None, None
            elif cds_union.size == 0:
                merged = utr_union
                cds_start, cds_end = None, None
            elif utr_union.size == 0:
                merged = cds_union
                cds_start = int(cds_union[:, 0].min())
                cds_end = int(cds_union[:, 1].max())
            else:
                merged = merge_intervals(np.vstack([cds_union, utr_union]))
                cds_start = int(cds_union[:, 0].min())
                cds_end = int(cds_union[:, 1].max())
        else:
            all_blocks = [t.blocks for t in ts if t.blocks.size > 0]
            if not all_blocks:
                raise ValueError(f"Gene {self.gene_id} has no exon blocks to merge")
            merged = merge_intervals(np.vstack(all_blocks))
            cds_start, cds_end = None, None

        return Transcript(
            blocks=merged,
            strand=self._strand,
            chrom=self._chrom,
            id=tid,
            cds_start=cds_start,
            cds_end=cds_end,
            biotype="exons_union",
        )

    def apply_canonizer(self, canonizer_func: Optional[Callable] = None, **kwargs) -> None:
        """Set the canonical transcript using a canonizer function.
        
        Args:
            canonizer_func: Function that takes transcripts list and returns canonical transcript ID.
                          If None, uses the default longest_isoform_canonizer.
            **kwargs: Additional arguments passed to the canonizer function.
        """
        if canonizer_func is None:
            from .canonizer import DEFAULT_CANONIZER
            canonizer_func = DEFAULT_CANONIZER
        
        canonical_id = canonizer_func(self._transcripts, **kwargs)
        
        if canonical_id is not None and canonical_id not in self.transcript_ids:
            raise ValueError(f"Canonical transcript ID '{canonical_id}' not found in gene {self.gene_id}")
        
        self._canonical_transcript_id = canonical_id


    def set_canonical_transcript(self, transcript_id: str) -> None:
        if transcript_id not in self.transcript_ids:
            raise ValueError(f"Transcript ID '{transcript_id}' not found in gene {self.gene_id}")
        
        self._canonical_transcript_id = transcript_id
    
    @property
    def canonical_transcript_id(self) -> Optional[str]:
        return self._canonical_transcript_id
    
    @property
    def canonical_transcript(self) -> Optional[Transcript]:
        if self._canonical_transcript_id is None:
            return None
        return self.get_transcript(self._canonical_transcript_id)
    
    @property
    def has_canonical_transcript(self) -> bool:
        return self._canonical_transcript_id is not None
    
    def clear_canonical_transcript(self) -> None:
        self._canonical_transcript_id = None
    
    def __len__(self) -> int:
        return len(self._transcripts)
    
    def __repr__(self) -> str:
        name_info = f", name={self.gene_name}" if self.gene_name else ""
        canonical_info = f", canonical={self._canonical_transcript_id}" if self.has_canonical_transcript else ""
        return f"Gene(id={self.gene_id}{name_info}, {self.chrom}:{self.start}-{self.end}:{self._strand}, {len(self._transcripts)} transcripts{canonical_info})"


class TranscriptsCollection:
    """Container for many transcripts."""

    def __init__(
        self,
        transcripts: Optional[List[Transcript]] = None,
        source_file: Optional[str] = None
    ):
        self.transcripts: List[Transcript] = transcripts or []
        self.source_file: Optional[str] = source_file

        self._id_index: Optional[Dict[str, int]] = None  # transcript_id → index
        self._chrom_index: Optional[Dict[str, List[int]]] = None  # chrom → indices
        self._gene_data: Optional['GeneData'] = None
        self._genes_cache: Optional[Dict[str, Gene]] = None

    def _invalidate_indices(self):
        self._id_index = None
        self._chrom_index = None
        self._genes_cache = None

    def append(self, transcript: Transcript) -> None:
        """Append a single transcript to the collection."""
        self.transcripts.append(transcript)
        self._invalidate_indices()

    def extend(self, other: Union[List[Transcript], 'TranscriptsCollection']) -> None:
        """Extend this collection with transcripts from a list or another collection."""
        if isinstance(other, TranscriptsCollection):
            self.transcripts.extend(other.transcripts)
        else:
            self.transcripts.extend(other)
        self._invalidate_indices()

    def __add__(self, other: 'TranscriptsCollection') -> 'TranscriptsCollection':
        """Return a new collection combining transcripts from both collections."""
        combined = TranscriptsCollection(
            transcripts=self.transcripts + other.transcripts,
            source_file=self.source_file,
        )
        if self._gene_data is not None:
            combined.bind_gene_data(self._gene_data)
        elif other._gene_data is not None:
            combined.bind_gene_data(other._gene_data)
        return combined

    def __iadd__(self, other: 'TranscriptsCollection') -> 'TranscriptsCollection':
        """In-place extend with another collection (+=)."""
        self.extend(other)
        return self

    def __len__(self):
        return len(self.transcripts)

    def __iter__(self):
        """Iterate over transcripts with bound metadata applied (biotype, etc.)."""
        for i in range(len(self.transcripts)):
            yield self._enrich_transcript(self.transcripts[i])

    def __getitem__(self, idx: int) -> Transcript:
        transcript = self.transcripts[idx]
        return self._enrich_transcript(transcript)

    def get_by_id(self, transcript_id: str) -> Optional[Transcript]:
        if self._id_index is None:
            self._build_id_index()
        idx = self._id_index.get(transcript_id)
        if idx is not None:
            transcript = self.transcripts[idx]
            return self._enrich_transcript(transcript)
        return None

    def get_by_chrom(self, chrom: str) -> List[Transcript]:
        if self._chrom_index is None:
            self._build_chrom_index()
        indices = self._chrom_index.get(chrom, [])
        return [self._enrich_transcript(self.transcripts[i]) for i in indices]

    def filter_by_chroms(self, chroms: Union[str, List[str], Set[str]]) -> 'TranscriptsCollection':
        """Return a new subcollection containing only transcripts on the given chromosome(s).

        Carries over bound GeneData and source_file from the parent collection.
        """
        if isinstance(chroms, str):
            chroms = {chroms}
        else:
            chroms = set(chroms)

        if self._chrom_index is None:
            self._build_chrom_index()

        indices = []
        for chrom in chroms:
            indices.extend(self._chrom_index.get(chrom, []))
        indices.sort()

        filtered = [self.transcripts[i] for i in indices]
        new_collection = TranscriptsCollection(transcripts=filtered, source_file=self.source_file)
        if self._gene_data is not None:
            new_collection.bind_gene_data(self._gene_data)
        return new_collection

    def filter_by_biotype(self, biotype: Union[str, List[str], Set[str]]) -> 'TranscriptsCollection':
        """Return a new collection containing only transcripts with the given biotype(s).

        Accepts a single biotype string or a collection of biotype strings.
        Uses bound GeneData for biotype when transcript.biotype is not set.
        """
        from .genes_auxiliary import filter_transcripts_by_biotype
        if isinstance(biotype, str):
            biotype = {biotype}
        else:
            biotype = set(biotype)
        return filter_transcripts_by_biotype(self, biotype)

    def get_by_biotype(self, biotype: Union[str, List[str], Set[str]]) -> 'TranscriptsCollection':
        """Alias for filter_by_biotype."""
        return self.filter_by_biotype(biotype)

    def get_all_chromosomes(self) -> List[str]:
        if self._chrom_index is None:
            self._build_chrom_index()
        return list(self._chrom_index.keys())

    def get_transcript_ids_by_chrom(self, chrom: str) -> List[str]:
        if self._chrom_index is None:
            self._build_chrom_index()
        indices = self._chrom_index.get(chrom, [])
        return [self.transcripts[i].id for i in indices]
    
    def get_transcripts_in_interval(self, interval: GenomicInterval, include_partial: bool = True) -> 'TranscriptsCollection':
        from .genes_auxiliary import filter_transcripts_in_interval
        return filter_transcripts_in_interval(self, interval, include_partial)

    def copy_with_remapped_ids(
        self,
        id_mapping: Union[Dict[str, str], Callable[[str], str]],
        source_file: Optional[str] = None,
    ) -> 'TranscriptsCollection':
        """Return a new collection with the same transcripts but different IDs.

        IDs are immutable on Transcript, so this creates new Transcript copies
        with the requested IDs. See pyrion.ops.transformations.remap_transcript_ids.
        """
        from ..ops.transformations import remap_transcript_ids
        return remap_transcript_ids(self, id_mapping, source_file=source_file)

    def _build_id_index(self):
        self._id_index = {t.id: i for i, t in enumerate(self.transcripts)}

    def _build_chrom_index(self):
        from collections import defaultdict
        chrom_index = defaultdict(list)
        for i, t in enumerate(self.transcripts):
            chrom_index[t.chrom].append(i)
        self._chrom_index = dict(chrom_index)

    def bind_gene_data(self, gene_data: 'GeneData') -> None:
        self._gene_data = gene_data
        self._genes_cache = None  # Clear cache when binding new data

    def _enrich_transcript(self, transcript: Transcript) -> Transcript:
        """Apply gene data to a single transcript (lazy evaluation)."""
        if self._gene_data is None:
            return transcript
        
        # Apply biotype if available and not already set
        biotype = transcript.biotype
        if biotype is None and self._gene_data.has_biotype_mapping():
            biotype = self._gene_data.get_transcript_biotype(transcript.id)
        
        # Only create new transcript if we have changes
        if biotype is not None and biotype != transcript.biotype:
            return Transcript(
                blocks=transcript.blocks,
                strand=transcript.strand,
                chrom=transcript.chrom,
                id=transcript.id,
                cds_start=transcript.cds_start,
                cds_end=transcript.cds_end,
                biotype=biotype
            )
        
        return transcript

    def _build_genes_cache(self):
        if not self.has_gene_mapping:
            raise ValueError("No gene data with gene-transcript mapping bound. Call bind_gene_data() first.")
        
        if self._id_index is None:
            self._build_id_index()
        
        genes_dict = {}
        
        for gene_id in self._gene_data.gene_ids:
            # Get transcripts for this gene that exist in our collection
            transcript_ids = self._gene_data.get_transcripts(gene_id)
            gene_transcripts = []
            
            for transcript_id in transcript_ids:
                transcript = self.get_by_id(transcript_id)
                if transcript is not None:
                    gene_transcripts.append(transcript)
            
            if gene_transcripts:
                gene_name = None
                if self._gene_data.has_gene_name_mapping():
                    gene_name = self._gene_data.get_gene_name(gene_id)
                
                genes_dict[gene_id] = Gene(gene_id, gene_transcripts, gene_name=gene_name)
        
        self._genes_cache = genes_dict

    def get_gene_by_id(self, gene_id: str) -> Optional[Gene]:
        if not self.has_gene_mapping:
            raise ValueError("No gene data with gene-transcript mapping bound. Call bind_gene_data() first.")
        
        if self._genes_cache is None:
            self._build_genes_cache()
        
        return self._genes_cache.get(gene_id)

    def get_gene_by_transcript_id(self, transcript_id: str) -> Optional[Gene]:
        if not self.has_gene_mapping:
            raise ValueError("No gene data with gene-transcript mapping bound. Call bind_gene_data() first.")
        
        gene_id = self._gene_data.get_gene(transcript_id)
        if gene_id is None:
            return None
        
        return self.get_gene_by_id(gene_id)

    def get_by_gene_name(self, gene_name: str) -> List[Gene]:
        """Get Gene objects by gene name. Multiple genes can have the same name."""
        if not self.has_gene_mapping:
            raise ValueError("No gene data with gene-transcript mapping bound. Call bind_gene_data() first.")
        
        if not self._gene_data.has_gene_name_mapping():
            raise ValueError("No gene name mapping available. Gene names were not provided in the gene data.")
        
        gene_ids = self._gene_data.get_genes_by_name(gene_name)
        
        genes = []
        for gene_id in gene_ids:
            gene = self.get_gene_by_id(gene_id)
            if gene is not None:
                genes.append(gene)
        
        return genes

    @property
    def genes(self) -> List[Gene]:
        if not self.has_gene_mapping:
            raise ValueError("No gene data with gene-transcript mapping bound. Call bind_gene_data() first.")
        
        if self._genes_cache is None:
            self._build_genes_cache()
        
        return list(self._genes_cache.values())

    @property
    def gene_ids(self) -> Set[str]:
        if not self.has_gene_mapping:
            raise ValueError("No gene data with gene-transcript mapping bound. Call bind_gene_data() first.")
        
        if self._genes_cache is None:
            self._build_genes_cache()
        
        return set(self._genes_cache.keys())

    @property
    def has_gene_mapping(self) -> bool:
        return self._gene_data is not None and self._gene_data.has_gene_transcript_mapping()
    
    @property
    def available_data_mappings(self) -> List[str]:
        if self._gene_data is None:
            return []
        
        mappings = []
        if self._gene_data.has_gene_transcript_mapping():
            mappings.append("gene-transcript")
        if self._gene_data.has_biotype_mapping():
            mappings.append("transcript-biotype")
        if self._gene_data.has_gene_name_mapping():
            mappings.append("gene-name")
        return mappings
    
    @property
    def applied_biotypes(self) -> bool:
        return getattr(self, '_applied_biotypes', False)
    
    @property
    def applied_gene_names(self) -> bool:
        return getattr(self, '_applied_gene_names', False)
    


    _SORT_KEYS = {
        "chrom": lambda t: t.chrom,
        "start": lambda t: t.start,
        "end": lambda t: t.end,
        "id": lambda t: t.id,
        "strand": lambda t: t.strand.value if isinstance(t.strand, Strand) else int(t.strand),
    }

    def sort(
        self,
        by: Union[str, List[str]] = "position",
        reverse: bool = False,
    ) -> 'TranscriptsCollection':
        """Sort transcripts in-place and return self for chaining.

        Args:
            by: Sort key(s). A single string or list of strings.
                Predefined presets:
                    "position" — sort by (chrom, start, end)  [default]
                    "id"       — sort by transcript ID
                Individual keys: "chrom", "start", "end", "id", "strand".
                Multiple keys: ["chrom", "start"] sorts by chrom first, then start.
            reverse: If True, sort in descending order.

        Returns:
            self (for chaining, e.g. ``collection.sort().save_to_bed12(path)``)
        """
        if isinstance(by, str):
            if by == "position":
                keys = ["chrom", "start", "end"]
            elif by in self._SORT_KEYS:
                keys = [by]
            else:
                raise ValueError(
                    f"Unknown sort key '{by}'. "
                    f"Use one of: {', '.join(sorted(self._SORT_KEYS))} or 'position'."
                )
        else:
            keys = list(by)
            for k in keys:
                if k not in self._SORT_KEYS:
                    raise ValueError(
                        f"Unknown sort key '{k}'. "
                        f"Use one of: {', '.join(sorted(self._SORT_KEYS))}."
                    )

        extractors = [self._SORT_KEYS[k] for k in keys]
        self.transcripts.sort(
            key=lambda t: tuple(fn(t) for fn in extractors),
            reverse=reverse,
        )
        self._invalidate_indices()
        return self

    def canonize_transcripts(self, canonizer_func: Optional[Callable] = None, **kwargs) -> None:
        from .genes_auxiliary import set_canonical_transcripts_for_collection
        set_canonical_transcripts_for_collection(self, canonizer_func, **kwargs)

    def get_canonical_transcripts(self) -> 'TranscriptsCollection':
        from .genes_auxiliary import get_canonical_transcripts_only_from_collection
        return get_canonical_transcripts_only_from_collection(self)

    def apply_gene_canonical_mapping(self, gene_to_canonical: Dict[str, str]) -> None:
        """gene_to_canonical: Dictionary mapping gene IDs to canonical transcript IDs"""
        if not self.has_gene_mapping:
            raise ValueError("No gene data with gene-transcript mapping bound. Call bind_gene_data() first.")
        
        # Ensure genes cache is built
        if self._genes_cache is None:
            self._build_genes_cache()
        
        # Apply mapping to each gene
        for gene_id, canonical_transcript_id in gene_to_canonical.items():
            gene = self._genes_cache.get(gene_id)
            if gene is not None:
                gene.set_canonical_transcript(canonical_transcript_id)

    def get_genes_without_canonical_transcript(self) -> List[Gene]:
        if not self.has_gene_mapping:
            raise ValueError("No gene data with gene-transcript mapping bound. Call bind_gene_data() first.")
        
        if self._genes_cache is None:
            self._build_genes_cache()
        
        return [gene for gene in self._genes_cache.values() if not gene.has_canonical_transcript]

    def summary(self) -> str:
        base_summary = f"{len(self.transcripts):,} transcripts"
        if self._gene_data is not None and self._gene_data.has_gene_transcript_mapping():
            if self._genes_cache is None:
                self._build_genes_cache()
            gene_count = len(self._genes_cache)
            base_summary += f", {gene_count:,} genes"
        
        source_str = f" from {self.source_file}" if self.source_file else ""
        return base_summary + source_str

    def __str__(self) -> str:
        from ..ops.transcript_serialization import transcripts_collection_summary_string
        return transcripts_collection_summary_string(self)
    
    def get_trimmed_gene_data(self) -> 'GeneData':
        """Return a copy of bound GeneData containing only entries for transcripts in this collection."""
        if self._gene_data is None:
            raise ValueError("No GeneData bound to this collection. Call bind_gene_data() first.")
        transcript_ids = {t.id for t in self.transcripts}
        return self._gene_data.subset(transcript_ids)

    def save_biodata(
        self, 
        tsv_path: Union[str, Path],
        include_gene_transcript: bool = True,
        include_transcript_biotype: bool = True,
        include_gene_name: bool = True,
        separator: str = '\t',
        trim: bool = True,
    ) -> None:
        if self._gene_data is None:
            raise ValueError("No GeneData bound to this collection. Call bind_gene_data() first or use read_gtf() to load data with gene information.")

        gene_data = self.get_trimmed_gene_data() if trim else self._gene_data

        from ..io.gene_data import write_gene_data_tsv
        write_gene_data_tsv(
            gene_data,
            tsv_path,
            include_gene_transcript=include_gene_transcript,
            include_transcript_biotype=include_transcript_biotype,
            include_gene_name=include_gene_name,
            separator=separator
        )

    def __repr__(self) -> str:
        return f"<TranscriptsCollection: {self.summary()}>"

    def to_bed12_string(self) -> str:
        from ..ops.transcript_serialization import transcripts_collection_to_bed12_string
        return transcripts_collection_to_bed12_string(self)

    def save_to_bed12(self, file_path: Union[str, Path]) -> None:
        from ..ops.transcript_serialization import save_transcripts_collection_to_bed12
        save_transcripts_collection_to_bed12(self, file_path)

    def save_to_json(self, file_path: Union[str, Path]) -> None:
        from ..ops.transcript_serialization import save_transcripts_collection_to_json
        save_transcripts_collection_to_json(self, file_path)

    @classmethod
    def from_json(cls, file_path: Union[str, Path]) -> 'TranscriptsCollection':
        from ..ops.transcript_serialization import load_transcripts_collection_from_json
        return load_transcripts_collection_from_json(file_path)

