from __future__ import annotations
import os,sys
from pathlib import Path
import torch

class StudentFlowRuntime:
    """Unified evaluation/runtime wrapper for v7 students before ONNX/TRT promotion.

    Loads the original AniGen flow once, merges a chain of compact LoRA stage adapters,
    then optionally applies velocity symmetry guidance and time-aware residual caching.
    The wrapper preserves the original AniGen model call signature, so existing samplers
    and pipelines can use it without changes.
    """
    def __init__(self,model,steps:int,cache_enabled:bool=True,cache_threshold=.985,cache_max_reuse=2,symmetry_group:str='none',symmetry_weight:float=0.0,symmetry_every:int=1,sparse_resolution:int=64):
        self.model=model;self.steps=int(steps);self.symmetry_group=symmetry_group;self.symmetry_weight=float(symmetry_weight);self.symmetry_every=max(1,int(symmetry_every));self.sparse_resolution=int(sparse_resolution);self.calls=0
        self.cached=None
        if cache_enabled and self.steps>=4:
            from .timestep_cache import CachedFlowModel
            self.cached=CachedFlowModel(model,threshold=cache_threshold,max_reuse=cache_max_reuse,max_timestep_gap=max(1.0,1100.0/self.steps))
    def reset(self):
        self.calls=0
        if self.cached:self.cached.reset()
    def _base(self,x,xs,t,cond,*args,**kwargs):
        return self.cached(x,xs,t,cond,*args,**kwargs) if self.cached else self.model(x,xs,t,cond,*args,**kwargs)
    def __call__(self,x,xs,t,cond,*args,**kwargs):
        self.calls+=1
        if self.symmetry_weight<=0 or self.symmetry_group in ('','none') or (self.calls-1)%self.symmetry_every:
            return self._base(x,xs,t,cond,*args,**kwargs)
        # Symmetry uses the raw underlying model for transformed states to avoid polluting
        # the residual cache with a different coordinate frame.
        from .symmetry import symmetrized_velocity
        v0,s0=self._base(x,xs,t,cond,*args,**kwargs)
        vg,sg=symmetrized_velocity(self.model,x,xs,t,cond,group=self.symmetry_group,res=self.sparse_resolution,*args,**kwargs)
        w=self.symmetry_weight
        if hasattr(v0,'replace'):
            return v0.replace(feats=v0.feats*(1-w)+vg.feats*w),s0.replace(feats=s0.feats*(1-w)+sg.feats*w)
        return v0*(1-w)+vg*w,s0*(1-w)+sg*w
    @property
    def metrics(self):
        return {'steps':self.steps,'calls':self.calls,'cache_hit_rate':self.cached.hit_rate if self.cached else 0.0,'symmetry_group':self.symmetry_group,'symmetry_weight':self.symmetry_weight}


def load_student(component:str,adapter_paths:list[str]|None=None,students_repo='patdev/Companion-Forge-v7-Students',token=None,app_root='/home/user/app',**runtime_kwargs):
    """Load base AniGen and merge an ordered adapter chain (10→4, 4→2, 2→1)."""
    # Import training utilities only for checkpoint/adapters; inference wrapper remains tiny.
    runtime_repo_root=Path(__file__).resolve().parents[1]
    if str(runtime_repo_root.parent) not in sys.path:sys.path.insert(0,str(runtime_repo_root.parent))
    from v7.training.distill_flow import load_anigen_model
    from v7.training.lora import apply_adapter
    from huggingface_hub import hf_hub_download
    model,_=load_anigen_model(component,root=app_root);model.eval()
    merged=[]
    for path in adapter_paths or []:
        local=hf_hub_download(students_repo,path,repo_type='model',token=token or os.environ.get('HF_TOKEN'))
        state=torch.load(local,map_location='cpu',weights_only=False);merged+=apply_adapter(model,state)
    steps=int(runtime_kwargs.pop('steps',1 if adapter_paths and any('2to1' in p for p in adapter_paths) else (2 if adapter_paths and any('4to2' in p for p in adapter_paths) else 4)))
    rt=StudentFlowRuntime(model,steps=steps,**runtime_kwargs);rt.merged_adapter_modules=merged;return rt


def aggregate_multiview_conditions(cond_dicts:list[dict],mode='mean'):
    """Shape-preserving multi-view fusion usable immediately with existing 1374-token engines.

    Mean fusion requires no learned weights and is therefore the safe fallback. A trained
    pose-aware MultiViewAggregator can replace it later while keeping the same output shape.
    """
    if not cond_dicts:raise ValueError('At least one view conditioning dict is required')
    if len(cond_dicts)==1:return cond_dicts[0]
    if mode!='mean':raise ValueError('Only weight-free mean fusion is production-safe before aggregator training')
    out={}
    keys=set().union(*(d.keys() for d in cond_dicts))
    for k in keys:
        vals=[d[k] for d in cond_dicts if k in d]
        if vals and torch.is_tensor(vals[0]) and vals[0].is_floating_point() and all(v.shape==vals[0].shape for v in vals):out[k]=torch.stack(vals).mean(0)
        else:out[k]=vals[0]
    return out