"""
Online Learning System
Continuously improves RL policy from agent experiences

INNOVATION: Learns from real-world experiences with fairness integration
- Experiences include fairness features in 25D state
- Rewards include fairness bonuses (low Gini, diversity)
- Training data preserves fairness objectives
"""
import asyncio
import pickle
import traceback
import numpy as np
import torch
from stable_baselines3 import PPO
from pathlib import Path
import time
from typing import Optional

from .experience_buffer import ExperienceBuffer
from ..config import RLConfig


class OnlineLearner:
    """
    Manages online learning for the RL policy
    
    Process:
    1. Agent collects experiences during operation
    2. Experiences stored in replay buffer
    3. Periodically, learner updates model from buffer
    4. New model deployed to agent
    """
    
    def __init__(self, node_id: str, config: RLConfig = None, policy=None, data_dir: str = "./data",
                 buffer_size: int = None, update_frequency: int = None):
        """
        Initialize OnlineLearner

        Args:
            node_id: Node identifier
            config: RLConfig instance (for production use)
            policy: RLPolicy instance (for production use)
            data_dir: Directory for storing data
            buffer_size: Buffer capacity (for test compatibility)
            update_frequency: Update interval in seconds (for test compatibility)
        """
        self.node_id = node_id
        self.policy = policy  # RLPolicy instance
        self.data_dir = Path(data_dir)

        # Support test-style initialization
        if buffer_size is not None:
            capacity = buffer_size
        else:
            capacity = 10000

        if update_frequency is not None:
            self.update_interval = update_frequency
        else:
            self.update_interval = 300  # Update every 5 minutes

        # Experience buffer
        self.buffer = ExperienceBuffer(capacity=capacity, data_dir=data_dir)

        # Learning config
        if config is not None:
            self.config = config
            self.learning_enabled = config.online_learning
        else:
            # Default for tests
            self.config = None
            self.learning_enabled = False

        self.min_buffer_size = 100  # Need 100 experiences before learning
        self.batch_size = 64

        # Training model (separate from inference model)
        self.training_model: Optional[PPO] = None

        # Exploration decay
        self.exploration_decay = getattr(config, 'exploration_decay', 0.995) if config else 0.995
        self.exploration_min = getattr(config, 'exploration_min', 0.01) if config else 0.01

        # Stats
        self.updates_performed = 0
        self.last_update_time = 0

        self.running = False
    
    async def start(self):
        """Start online learning loop"""
        if not self.learning_enabled:
            print("[ONLINE LEARNER] Online learning disabled")
            return
        
        self.running = True
        
        # Load or create training model
        self._initialize_training_model()
        
        # Start learning loop
        asyncio.create_task(self._learning_loop())
        
        print(f"[ONLINE LEARNER] Started (update_interval={self.update_interval}s)")
    
    async def stop(self):
        """Stop online learning"""
        self.running = False
        
        # Save final model
        if self.training_model:
            self._save_training_model()
    
    def record_experience(self, state: np.ndarray, action: int, reward: float,
                         next_state: np.ndarray, done: bool):
        """Record experience from agent"""
        self.buffer.add(state, action, reward, next_state, done)
    
    async def _learning_loop(self):
        """Continuous learning loop"""
        while self.running:
            await asyncio.sleep(self.update_interval)
            
            # Check if we have enough experiences
            if self.buffer.size() < self.min_buffer_size:
                print(f"[ONLINE LEARNER] Buffer too small ({self.buffer.size()}/{self.min_buffer_size}), skipping update")
                continue
            
            # Perform update
            await self._perform_update()
    
    async def _perform_update(self):
        """Perform model update from buffer"""
        print(f"\n{'='*60}")
        print(f"[ONLINE LEARNER] Starting model update #{self.updates_performed + 1}")
        print(f"{'='*60}")
        
        start_time = time.time()
        
        try:
            # Get experiences from buffer
            experiences = self.buffer.get_all()
            
            print(f"[ONLINE LEARNER] Training on {len(experiences)} experiences")
            
            # Convert to training format
            states = np.array([exp.state for exp in experiences])
            actions = np.array([exp.action for exp in experiences])
            rewards = np.array([exp.reward for exp in experiences])
            next_states = np.array([exp.next_state for exp in experiences])
            dones = np.array([exp.done for exp in experiences])
            
            # Statistics before training
            stats_before = self.buffer.get_statistics()
            print(f"[ONLINE LEARNER] Avg reward before: {stats_before['avg_reward']:.3f}")
            print(f"[ONLINE LEARNER] Success rate: {stats_before['success_rate']:.2%}")
            
            # Perform training update
            
            # Behavioral cloning: imitate successful actions
            # Filter to experiences with positive reward (good decisions to reinforce)
            successful = [e for e in experiences if e.reward > 0]

            if len(successful) < self.batch_size:
                print(f"[ONLINE LEARNER] Not enough successful experiences ({len(successful)}), skipping BC update")
            elif self.training_model is None:
                print(f"[ONLINE LEARNER] No training model available, skipping update")
            else:
                num_updates = min(10, len(successful) // self.batch_size)
                total_loss = 0.0

                for i in range(num_updates):
                    batch_experiences = self.buffer.sample(self.batch_size)
                    good_batch = [e for e in batch_experiences if e.reward > 0]
                    if not good_batch:
                        continue

                    states = np.array([e.state for e in good_batch])
                    actions = np.array([e.action for e in good_batch])

                    try:
                        from stable_baselines3.common.utils import obs_as_tensor
                        obs_t = obs_as_tensor(states, self.training_model.device)
                        act_t = torch.tensor(actions, dtype=torch.long,
                                             device=self.training_model.device)

                        # Get action log-probs from current policy
                        distribution = self.training_model.policy.get_distribution(obs_t)
                        log_probs = distribution.log_prob(act_t)

                        # Behavioral cloning loss: maximize log-prob of good actions
                        loss = -log_probs.mean()
                        total_loss += loss.item()

                        self.training_model.policy.optimizer.zero_grad()
                        loss.backward()
                        # Gradient clipping to avoid instability
                        torch.nn.utils.clip_grad_norm_(
                            self.training_model.policy.parameters(), max_norm=0.5
                        )
                        self.training_model.policy.optimizer.step()

                    except Exception as update_err:
                        print(f"[ONLINE LEARNER] BC update error (step {i}): {update_err}")
                        break

                if num_updates > 0:
                    avg_loss = total_loss / num_updates
                    print(f"[ONLINE LEARNER] BC updates: {num_updates}, avg loss: {avg_loss:.4f}")
            
            # Simpler approach: Retrain model periodically
            if len(experiences) >= 500 and self.updates_performed % 5 == 0:
                print("[ONLINE LEARNER] Performing full retraining...")
                await self._retrain_model(experiences)
            
            # Decay exploration rate
            if self.policy and hasattr(self.policy, 'exploration_rate'):
                old_rate = self.policy.exploration_rate
                self.policy.exploration_rate = max(
                    self.exploration_min,
                    self.policy.exploration_rate * self.exploration_decay
                )
                print(f"[ONLINE LEARNER] Exploration: {old_rate:.4f} -> {self.policy.exploration_rate:.4f}")

            duration = time.time() - start_time
            self.updates_performed += 1
            self.last_update_time = time.time()

            print(f"[ONLINE LEARNER] Update completed in {duration:.2f}s")
            print(f"{'='*60}\n")
        
        except Exception as e:
            print(f"[ONLINE LEARNER] Error during update: {e}")
            traceback.print_exc()
    
    async def _retrain_model(self, experiences: list):
        """
        Save accumulated experiences for offline retraining.
        Called every 5th update cycle once 500+ experiences are collected.
        """
        self._save_experiences_for_offline_training(experiences)
        print(f"[ONLINE LEARNER] {len(experiences)} experiences saved for offline retraining at {self.data_dir}/experiences_for_training.pkl")
    
    def _save_experiences_for_offline_training(self, experiences: list):
        """Save experiences for offline retraining"""
        experience_file = self.data_dir / "experiences_for_training.pkl"
        
        with open(experience_file, 'wb') as f:
            pickle.dump(experiences, f)
        
        print(f"[ONLINE LEARNER] Saved {len(experiences)} experiences to {experience_file}")
    
    def _initialize_training_model(self):
        """Initialize or load training model"""
        training_model_path = self.data_dir / "training_model.zip"
        
        if training_model_path.exists():
            try:
                self.training_model = PPO.load(str(training_model_path))
                print(f"[ONLINE LEARNER] Loaded training model from {training_model_path}")
            except Exception as e:
                print(f"[ONLINE LEARNER] Error loading training model: {e}")
                self.training_model = None
        else:
            # Copy from inference model
            if self.policy and self.policy.model:
                self.training_model = self.policy.model
                print("[ONLINE LEARNER] Using inference model for training")
    
    def _save_training_model(self):
        """Save training model to disk"""
        if self.training_model:
            training_model_path = self.data_dir / "training_model.zip"
            self.training_model.save(str(training_model_path))
            print(f"[ONLINE LEARNER] Saved training model to {training_model_path}")
    
    def get_learning_stats(self) -> dict:
        """Get learning statistics"""
        buffer_stats = self.buffer.get_statistics()

        return {
            'learning_enabled': self.learning_enabled,
            'updates_performed': self.updates_performed,
            'last_update': self.last_update_time,
            'buffer_size': self.buffer.size(),
            'buffer_stats': buffer_stats,
            'next_update_in': max(0, self.update_interval - (time.time() - self.last_update_time))
        }

    def get_stats(self) -> dict:
        """
        Get statistics (alias for compatibility with tests)

        Returns:
            Dictionary with total_experiences and buffer_size
        """
        return {
            'total_experiences': self.buffer.size(),
            'buffer_size': self.buffer.size(),
            'updates_performed': self.updates_performed
        }

    def export_for_training(self, output_path: str) -> int:
        """
        Export experiences to file for offline training (compatibility method for tests)

        Args:
            output_path: Path to save experiences

        Returns:
            Number of experiences exported
        """
        experiences = self.buffer.get_all()

        output_file = Path(output_path)
        output_file.parent.mkdir(parents=True, exist_ok=True)

        with open(output_file, 'wb') as f:
            pickle.dump(experiences, f)

        print(f"[ONLINE LEARNER] Exported {len(experiences)} experiences to {output_path}")

        return len(experiences)