"""
MoE Expert Offloading - Move cold experts to CPU.

Keeps only active (hot) experts on GPU to save VRAM.
Critical for running large MoE models on limited GPU memory (e.g., 2x T4).
"""

import gc 
from typing import Dict ,List ,Optional ,Set 
from collections import defaultdict 
import threading 

import torch 
import torch .nn as nn 


class ExpertOffloadManager :
    """
    Manages CPU offloading of cold MoE experts.
    
    Usage:
        manager = ExpertOffloadManager(moe_layer.experts)
        
        
        manager.prefetch(selected_expert_indices)
        output = moe_forward(...)
        manager.record_usage(selected_expert_indices)
        
        
        manager.offload_cold()
    """

    def __init__ (
    self ,
    experts :nn .ModuleList ,
    device :str ='cuda:0',
    cold_threshold :int =10 ,
    max_offload_ratio :float =0.5 ,
    ):
        self .experts =experts 
        self .num_experts =len (experts )
        self .device =torch .device (device )
        self .cpu =torch .device ('cpu')
        self .cold_threshold =cold_threshold 
        self .max_offload =int (self .num_experts *max_offload_ratio )


        self .locations :Dict [int ,str ]={i :'gpu'for i in range (self .num_experts )}
        self .usage_count :Dict [int ,int ]=defaultdict (int )
        self .last_used :Dict [int ,int ]={}
        self .step =0 


        self .cpu_states :Dict [int ,Dict [str ,torch .Tensor ]]={}


        self ._lock =threading .Lock ()


        self ._stream =torch .cuda .Stream ()if torch .cuda .is_available ()else None 

    def prefetch (self ,expert_indices :torch .Tensor ):
        """Prefetch needed experts from CPU to GPU."""
        needed =set (expert_indices .unique ().tolist ())
        offloaded ={i for i in needed if self .locations .get (i )=='cpu'}

        if not offloaded :
            return 

        with self ._lock :
            for idx in offloaded :
                self ._load_to_gpu (idx )

    def record_usage (self ,expert_indices :torch .Tensor ):
        """Record which experts were used."""
        for idx in expert_indices .unique ().tolist ():
            self .usage_count [idx ]+=1 
            self .last_used [idx ]=self .step 

    def step_forward (self ):
        """Advance step counter."""
        self .step +=1 

    def offload_cold (self )->int :
        """Offload cold experts to CPU. Returns count offloaded."""
        cold =self ._get_cold_experts ()
        cold .sort (key =lambda i :self .last_used .get (i ,-999 ))

        count =0 
        with self ._lock :
            for idx in cold [:self .max_offload ]:
                if self .locations .get (idx )=='gpu':
                    self ._offload_to_cpu (idx )
                    count +=1 

        if count :
            gc .collect ()
            torch .cuda .empty_cache ()

        return count 

    def _get_cold_experts (self )->List [int ]:
        """Get indices of cold (unused) experts."""
        cold =[]
        for i in range (self .num_experts ):
            last =self .last_used .get (i ,-self .cold_threshold -1 )
            if self .step -last >self .cold_threshold :
                cold .append (i )
        return cold 

    def _load_to_gpu (self ,idx :int ):
        """Load expert from CPU to GPU."""
        if idx not in self .cpu_states :
            return 

        expert =self .experts [idx ]
        state =self .cpu_states [idx ]

        if self ._stream :
            with torch .cuda .stream (self ._stream ):
                for name ,param in expert .named_parameters ():
                    if name in state :
                        param .data =state [name ].to (self .device ,non_blocking =True )
            self ._stream .synchronize ()
        else :
            for name ,param in expert .named_parameters ():
                if name in state :
                    param .data =state [name ].to (self .device )

        self .locations [idx ]='gpu'
        del self .cpu_states [idx ]

    def _offload_to_cpu (self ,idx :int ):
        """Offload expert from GPU to CPU."""
        expert =self .experts [idx ]

        state ={}
        for name ,param in expert .named_parameters ():
            state [name ]=param .data .to (self .cpu ,non_blocking =True )

        self .cpu_states [idx ]=state 


        for param in expert .parameters ():
            param .data =torch .empty (0 ,device =self .cpu )

        self .locations [idx ]='cpu'

    def load_all_to_gpu (self ):
        """Load all experts back to GPU."""
        with self ._lock :
            for idx in list (self .cpu_states .keys ()):
                self ._load_to_gpu (idx )

    def get_status (self )->Dict :
        """Get current status."""
        gpu =sum (1 for v in self .locations .values ()if v =='gpu')
        cpu =sum (1 for v in self .locations .values ()if v =='cpu')
        return {
        'total':self .num_experts ,
        'gpu':gpu ,'cpu':cpu ,
        'step':self .step ,
        }

    def print_status (self ):
        """Print human-readable status."""
        s =self .get_status ()
        print (f"  Expert Offload: GPU={s ['gpu']}/{s ['total']}, CPU={s ['cpu']}/{s ['total']}")


def apply_expert_offloading (
model :nn .Module ,
cold_threshold :int =10 ,
max_offload_ratio :float =0.5 ,
)->Optional [ExpertOffloadManager ]:
    """
    Apply expert offloading to a model with MoE layers.
    
    Returns ExpertOffloadManager if MoE layers found, else None.
    """

    for name ,module in model .named_modules ():
        if hasattr (module ,'experts')and isinstance (module .experts ,nn .ModuleList ):
            manager =ExpertOffloadManager (
            module .experts ,
            cold_threshold =cold_threshold ,
            max_offload_ratio =max_offload_ratio ,
            )
            module ._expert_offload_manager =manager 
            print (f"  ✅ Expert offloading enabled for {name }")
            return manager 

    return None 
