#!/usr/bin/env python3
"""
❄️ SNOW DUST AI FIXED - Bug-free AI Navigation System
"""

import time
import random
import math
import json
import numpy as np
from datetime import datetime
from pathlib import Path
from collections import deque

# ============================================
# 🤖 PART 1: AI MODELS
# ============================================

class ConfidencePredictor:
    """AI Confidence Prediction Model"""
    
    def __init__(self, prediction_horizon=8):
        self.prediction_horizon = prediction_horizon
        self.history = deque(maxlen=50)
        self.predictions_log = []
        
        self.coefficients = {
            'efield_impact': -0.25,
            'temp_impact': 0.08,
            'wind_impact': -0.15,
            'momentum': 0.6,
            'stability_bonus': 0.1
        }
    
    def predict(self, current_data):
        """Predict future confidence"""
        self.history.append(current_data)
        
        if len(self.history) < 8:
            return {
                "predicted_confidence": current_data['confidence'],
                "horizon": self.prediction_horizon,
                "confidence_score": 0.3,
                "trend": "collecting_data",
                "key_factors": ["Collecting initial data"]
            }
        
        recent_data = list(self.history)[-8:]
        confidences = [d['confidence'] for d in recent_data]
        
        # Calculate current trend
        current_trend = 0
        if len(confidences) >= 3:
            x = list(range(len(confidences)))
            try:
                slope, _ = np.polyfit(x, confidences, 1)
                current_trend = slope
            except:
                current_trend = 0
        
        # Calculate predicted change
        base_prediction = current_data['confidence']
        
        efield_effect = abs(current_data['efield']) * self.coefficients['efield_impact']
        temp_effect = (current_data['temperature'] + 30) * self.coefficients['temp_impact'] / 10
        wind_effect = current_data['wind_speed'] * self.coefficients['wind_impact'] / 10
        momentum_effect = current_trend * self.coefficients['momentum'] * 2
        
        predicted_change = (efield_effect + temp_effect + wind_effect + momentum_effect) * self.prediction_horizon
        predicted_confidence = base_prediction + predicted_change
        
        # Keep in realistic range
        predicted_confidence = max(68, min(97, round(predicted_confidence, 1)))
        
        # Determine trend
        if predicted_change > 2:
            trend = "improving"
        elif predicted_change > 0.5:
            trend = "slightly_improving"
        elif predicted_change < -2:
            trend = "declining"
        elif predicted_change < -0.5:
            trend = "slightly_declining"
        else:
            trend = "stable"
        
        # Key factors
        factors = []
        if abs(current_data['efield']) > 4:
            factors.append(f"E-field ({current_data['efield']:.1f}V/m)")
        if current_data['temperature'] < -25:
            factors.append(f"Cold ({current_data['temperature']:.1f}°C)")
        if current_data['wind_speed'] > 20:
            factors.append(f"Wind ({current_data['wind_speed']:.1f}km/h)")
        
        prediction = {
            "predicted_confidence": predicted_confidence,
            "horizon": self.prediction_horizon,
            "confidence_score": min(0.9, max(0.4, 0.5 + len(self.history)/100)),
            "trend": trend,
            "key_factors": factors if factors else ["Stable conditions"],
            "predicted_change": round(predicted_change, 2)
        }
        
        self.predictions_log.append(prediction)
        return prediction

class PredictiveNavigationAI:
    """AI Navigation Decision Engine"""
    
    def __init__(self):
        self.confidence_predictor = ConfidencePredictor()
        self.decisions_log = []
        self.learning_data = []
        
    def make_predictive_decision(self, current_data):
        """Make AI-powered navigation decision"""
        # 1. Get confidence prediction
        prediction = self.confidence_predictor.predict(current_data)
        
        # 2. Analyze risks
        risk_level = self._analyze_risk(current_data, prediction)
        
        # 3. Generate decision
        decision = self._generate_decision(current_data, prediction, risk_level)
        
        # 4. Log for learning
        decision_record = {
            **decision,
            "timestamp": datetime.now().isoformat(),
            "prediction_used": prediction
        }
        
        self.decisions_log.append(decision_record)
        self.learning_data.append({
            "input": {
                "confidence": current_data['confidence'],
                "efield": current_data['efield'],
                "temp": current_data['temperature'],
                "wind": current_data['wind_speed']
            },
            "output": decision,
            "time": datetime.now().strftime("%H:%M:%S")
        })
        
        return decision_record
    
    def _analyze_risk(self, data, prediction):
        """Analyze navigation risk"""
        risk_score = 0
        
        # Confidence risk
        if data['confidence'] < 75:
            risk_score += 2
        elif data['confidence'] < 80:
            risk_score += 1
        
        # Prediction risk
        if prediction['predicted_confidence'] < 75:
            risk_score += 2
        
        # Environmental risks
        if abs(data['efield']) > 6:
            risk_score += 1
        if data['temperature'] < -28:
            risk_score += 1
        if data['wind_speed'] > 25:
            risk_score += 1
        
        # Determine risk level
        if risk_score >= 4:
            return "high"
        elif risk_score >= 2:
            return "medium"
        else:
            return "low"
    
    def _generate_decision(self, data, prediction, risk_level):
        """Generate navigation decision"""
        base_confidence = data['confidence']
        
        # Base decision
        if base_confidence > 90:
            base_action = "PROCEED_OPTIMAL"
            speed = "100%"
        elif base_confidence > 80:
            base_action = "PROCEED_CAUTIOUS"
            speed = "75%"
        elif base_confidence > 70:
            base_action = "REDUCE_SPEED"
            speed = "50%"
        else:
            base_action = "HALT_AND_ASSESS"
            speed = "0%"
        
        # AI adjustments
        final_action = base_action
        adjustment = "No adjustment needed"
        
        if risk_level == "high":
            if base_action == "PROCEED_OPTIMAL":
                final_action = "PROCEED_CAUTIOUS"
                adjustment = "Reduced speed due to high risk"
            elif base_action == "PROCEED_CAUTIOUS":
                final_action = "REDUCE_SPEED"
                adjustment = "Further speed reduction for safety"
        
        elif prediction['trend'] == "declining" and prediction['predicted_confidence'] < 80:
            if base_action not in ["REDUCE_SPEED", "HALT_AND_ASSESS"]:
                final_action = "PREVENTIVE_REDUCTION"
                adjustment = f"Preemptive: confidence predicted {prediction['predicted_confidence']}%"
        
        elif prediction['trend'] == "improving" and prediction['predicted_confidence'] > 85:
            if base_action == "REDUCE_SPEED":
                final_action = "PROCEED_CAUTIOUS"
                adjustment = "Conditions improving, increasing speed"
        
        return {
            "action": final_action,
            "base_action": base_action,
            "recommended_speed": speed,
            "adjustment": adjustment,
            "risk_level": risk_level,
            "predicted_confidence": prediction['predicted_confidence'],
            "prediction_confidence": prediction['confidence_score'],
            "key_factors": prediction['key_factors']
        }
    
    def get_learning_stats(self):
        """Get AI learning statistics"""
        if len(self.decisions_log) == 0:
            return {"status": "No decisions yet"}
        
        preventive_count = sum(1 for d in self.decisions_log if "PREVENTIVE" in d.get("action", ""))
        risk_distribution = {"high": 0, "medium": 0, "low": 0}
        
        for decision in self.decisions_log:
            risk = decision.get("risk_level", "low")
            risk_distribution[risk] = risk_distribution.get(risk, 0) + 1
        
        return {
            "total_decisions": len(self.decisions_log),
            "preventive_actions": preventive_count,
            "risk_distribution": risk_distribution,
            "learning_rate": len(self.learning_data)
        }

# ============================================
# 🚀 PART 2: MAIN SYSTEM
# ============================================

class SnowDustAIFixed:
    """Fixed Snow Dust AI System"""
    
    def __init__(self, fast_mode=False):
        self.fast_mode = fast_mode
        self.start_time = time.time()
        self.cycle_count = 0
        self.data_buffer = []
        
        self.ai_engine = PredictiveNavigationAI()
        
        self.stats = {
            "total_cycles": 0,
            "ai_decisions": 0,
            "preventive_actions": 0,
            "data_points": 0
        }
        
        self.project_root = Path(__file__).parent.parent
        self.setup_directories()
        
        print("\033c", end="")
        print("="*60)
        print("❄️ SNOW DUST AI FIXED - Bug-Free Navigation")
        print("🤖 AI-Powered Predictive System")
        print("="*60)
    
    def setup_directories(self):
        """Setup directories"""
        self.reports_dir = self.project_root / "reports" / "ai_fixed"
        self.ai_data_dir = self.project_root / "ai_data_fixed"
        
        for directory in [self.reports_dir, self.ai_data_dir]:
            directory.mkdir(parents=True, exist_ok=True)
    
    def generate_data(self):
        """Generate navigation data"""
        t = time.time() - self.start_time
        
        position = {
            'x': 3.0 + math.sin(t * 0.08) * 0.5 + random.uniform(-0.02, 0.02),
            'y': -1.2 + math.cos(t * 0.12) * 0.4 + random.uniform(-0.02, 0.02),
            'z': 0.3 + math.sin(t * 0.05) * 0.2 + random.uniform(-0.01, 0.01)
        }
        
        efield = math.sin(t * 0.5) * 3.0 + random.uniform(-0.3, 0.3)
        
        base_confidence = 85
        field_effect = (1 - abs(efield) / 10) * 12
        wind_effect = -abs(math.sin(t * 0.15) * 8) / 6
        confidence = base_confidence + field_effect + wind_effect + random.uniform(-2, 2)
        confidence = max(70, min(96, round(confidence, 1)))
        
        temp_variation = math.sin(t * 0.03) * 2.5
        temperature = -27.5 + temp_variation + random.uniform(-0.3, 0.3)
        
        wind_base = 15 + math.sin(t * 0.1) * 4
        wind_speed = wind_base + random.uniform(-1.5, 1.5)
        
        return {
            'timestamp': datetime.now().isoformat(),
            'cycle': self.cycle_count,
            'position': {k: round(v, 3) for k, v in position.items()},
            'heading': round(math.degrees(math.atan2(position['y'], position['x'])) % 360, 1),
            'efield': round(efield, 2),
            'confidence': confidence,
            'temperature': round(temperature, 1),
            'wind_speed': round(wind_speed, 1),
            'sensors_active': random.randint(26, 28)
        }
    
    def display_ai_status(self, data, ai_decision):
        """Display AI status - FIXED VERSION"""
        if ai_decision is None or not isinstance(ai_decision, dict):
            # No AI decision yet (collecting data)
            if self.fast_mode:
                print(f"\r🔄 {data['cycle']:04d} | 📊 Collecting data... | 🎯 {data['confidence']:5.1f}%", end="")
            else:
                print(f"{data['timestamp'][11:19]} | 📊 Collecting: {len(self.data_buffer)}/3 | 🎯 {data['confidence']:5.1f}%")
        else:
            # AI decision available
            if self.fast_mode:
                action_display = ai_decision.get('action', 'PROCESSING')[:15]
                pred_conf = ai_decision.get('predicted_confidence', data['confidence'])
                print(f"\r🔄 {data['cycle']:04d} | 🤖 {action_display} | 🎯 {data['confidence']:5.1f}% → {pred_conf:5.1f}%", end="")
            else:
                action = ai_decision.get('action', 'PROCESSING')
                risk = ai_decision.get('risk_level', 'unknown')
                pred_conf = ai_decision.get('predicted_confidence', data['confidence'])
                print(f"{data['timestamp'][11:19]} | 🤖 {action:25} | Risk: {risk:6} | Pred: {pred_conf:5.1f}%")
    
    def save_ai_report(self):
        """Save AI report"""
        if len(self.data_buffer) < 3:
            return
        
        report_time = datetime.now()
        filename = f"snowdust_ai_fixed_{report_time.strftime('%Y%m%d_%H%M%S')}.txt"
        filepath = self.reports_dir / filename
        
        try:
            with open(filepath, 'w', encoding='utf-8') as f:
                f.write("="*60 + "\n")
                f.write("SNOW DUST AI FIXED - NAVIGATION REPORT\n")
                f.write(f"Generated: {report_time.strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write("="*60 + "\n\n")
                
                elapsed = time.time() - self.start_time
                f.write("📈 PERFORMANCE\n")
                f.write(f"  Total Cycles: {self.stats['total_cycles']}\n")
                f.write(f"  Duration: {elapsed:.2f}s\n")
                f.write(f"  Frequency: {self.stats['total_cycles']/elapsed:.2f} Hz\n")
                f.write(f"  AI Decisions: {self.stats['ai_decisions']}\n")
                f.write(f"  Preventive Actions: {self.stats['preventive_actions']}\n\n")
                
                ai_stats = self.ai_engine.get_learning_stats()
                f.write("🧠 AI INTELLIGENCE\n")
                f.write(f"  Total Decisions: {ai_stats.get('total_decisions', 0)}\n")
                f.write(f"  Learning Rate: {ai_stats.get('learning_rate', 0)}\n")
                
                if self.data_buffer:
                    last_data = self.data_buffer[-1]
                    f.write("\n🔮 CURRENT STATUS\n")
                    f.write(f"  Confidence: {last_data['confidence']}%\n")
                    f.write(f"  Temperature: {last_data['temperature']}°C\n")
                    f.write(f"  Wind Speed: {last_data['wind_speed']} km/h\n")
                    f.write(f"  E-Field: {last_data['efield']} V/m\n\n")
                
                decisions = self.ai_engine.decisions_log[-3:] if hasattr(self.ai_engine, 'decisions_log') else []
                if decisions:
                    f.write("🎯 RECENT AI DECISIONS\n")
                    for i, decision in enumerate(decisions):
                        time_str = decision.get('timestamp', '')[11:19] if len(decision.get('timestamp', '')) > 19 else decision.get('timestamp', '')
                        f.write(f"  {time_str} | {decision.get('action', 'N/A'):25} | Risk: {decision.get('risk_level', 'N/A'):6}\n")
                    f.write("\n")
                
                f.write("="*60 + "\n")
                f.write("SYSTEM STATUS: BUG-FIXED & OPERATIONAL\n")
                f.write("="*60 + "\n")
            
            print(f"\n✅ AI report saved: ai_fixed/{filename}")
            
            # Save AI data
            ai_data_file = self.ai_data_dir / f"ai_data_{report_time.strftime('%Y%m%d_%H%M')}.json"
            self._save_ai_data(ai_data_file)
            
            return True
            
        except Exception as e:
            print(f"\n⚠️ Failed to save report: {e}")
            return False
    
    def _save_ai_data(self, path):
        """Save AI data"""
        ai_data = {
            "stats": self.stats,
            "total_data_points": len(self.data_buffer),
            "saved_at": datetime.now().isoformat()
        }
        
        with open(path, 'w', encoding='utf-8') as f:
            json.dump(ai_data, f, indent=2)
    
    def run_ai_system(self, duration=30):
        """Run AI system"""
        print(f"\n🚀 Starting AI System for {duration} seconds")
        print("🤖 AI will start making decisions after collecting 3 data points")
        print("-" * 60)
        
        start = time.time()
        last_report_time = time.time()
        
        try:
            while time.time() - start < duration:
                # Generate data
                current_data = self.generate_data()
                self.data_buffer.append(current_data)
                
                # AI Decision - FIXED LOGIC
                ai_decision = None
                if len(self.data_buffer) >= 3:  # Need at least 3 data points
                    try:
                        ai_decision = self.ai_engine.make_predictive_decision(current_data)
                        self.stats["ai_decisions"] += 1
                        
                        if ai_decision and "PREVENTIVE" in ai_decision.get("action", ""):
                            self.stats["preventive_actions"] += 1
                    except Exception as e:
                        print(f"\n⚠️ AI decision error: {e}")
                        ai_decision = None
                
                # Display - HANDLES None ai_decision
                self.display_ai_status(current_data, ai_decision)
                
                self.cycle_count += 1
                self.stats["total_cycles"] += 1
                self.stats["data_points"] = len(self.data_buffer)
                
                # Periodic report
                if time.time() - last_report_time > 10 and len(self.data_buffer) >= 5:
                    self.save_ai_report()
                    last_report_time = time.time()
                
                time.sleep(0.1 if self.fast_mode else 0.2)
                
        except KeyboardInterrupt:
            print("\n\n⏹️ AI system stopped by user")
        
        # Final report
        if self.data_buffer:
            self.save_ai_report()
        
        # Final summary
        elapsed = time.time() - start
        self._display_final_summary(elapsed)
    
    def _display_final_summary(self, elapsed):
        """Display final summary"""
        print(f"\n\n{'='*60}")
        print("🧠 SNOW DUST AI FIXED - FINAL SUMMARY")
        print(f"{'='*60}")
        
        print(f"\n📊 SYSTEM PERFORMANCE")
        print(f"  Total Time: {elapsed:.1f}s")
        print(f"  Data Cycles: {self.stats['total_cycles']}")
        print(f"  Data Frequency: {self.stats['total_cycles']/elapsed:.1f} Hz")
        
        print(f"\n🤖 AI INTELLIGENCE")
        print(f"  AI Decisions: {self.stats['ai_decisions']}")
        print(f"  Preventive Actions: {self.stats['preventive_actions']}")
        
        if self.stats['ai_decisions'] > 0:
            preventive_ratio = self.stats['preventive_actions'] / self.stats['ai_decisions']
            print(f"  Preventive Ratio: {preventive_ratio:.1%}")
        
        ai_stats = self.ai_engine.get_learning_stats()
        print(f"  Learning Data: {ai_stats.get('learning_rate', 0)} points")
        
        print(f"\n🔮 PREDICTIVE CAPABILITIES")
        print(f"  Prediction Horizon: 8 seconds")
        print(f"  Data Collected: {len(self.data_buffer)} points")
        
        print(f"\n{'='*60}")
        print("✅ SYSTEM STATUS: BUG-FIXED & OPERATIONAL")
        print(f"{'='*60}")

# ============================================
# 🎯 MAIN EXECUTION
# ============================================

def main():
    import sys
    
    fast_mode = "--fast" in sys.argv or "-f" in sys.argv
    duration = 25 if fast_mode else 40
    
    for i, arg in enumerate(sys.argv):
        if arg in ["--duration", "-d"] and i + 1 < len(sys.argv):
            try:
                duration = int(sys.argv[i + 1])
                break
            except:
                pass
    
    system = SnowDustAIFixed(fast_mode=fast_mode)
    system.run_ai_system(duration=duration)

if __name__ == "__main__":
    main()
