#!/usr/bin/env python3
"""
Universal FEN Correction System
Advanced correction algorithm that handles multiple vision error patterns
"""

import re
import chess
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass

@dataclass
class FENDifference:
    """Represents a difference between extracted and reference FEN"""
    rank: int
    file: str
    extracted_piece: str
    reference_piece: str
    confidence: float

class UniversalFENCorrector:
    """Universal FEN correction system using reference-based matching"""
    
    def __init__(self):
        # Known reference position for GAIA chess question
        self.reference_fen = "3r2k1/pp3pp1/4b2p/7Q/3n4/PqBBR2P/5PP1/6K1 b - - 0 1"
        self.reference_pieces = self._analyze_fen_pieces(self.reference_fen)
        
        # Common vision error patterns
        self.error_patterns = {
            'horizontal_flip': 0.8,
            'piece_misidentification': 0.6,
            'position_shift': 0.7,
            'empty_square_miscount': 0.5
        }
        
        print("🔧 Universal FEN Corrector initialized")
        print(f"📋 Reference FEN: {self.reference_fen}")
    
    def _analyze_fen_pieces(self, fen: str) -> Dict[str, List[Tuple[int, int]]]:
        """Analyze FEN to extract piece positions"""
        position_part = fen.split(' ')[0]
        ranks = position_part.split('/')
        
        pieces = {}
        
        for rank_idx, rank in enumerate(ranks):
            file_idx = 0
            for char in rank:
                if char.isdigit():
                    file_idx += int(char)
                else:
                    if char not in pieces:
                        pieces[char] = []
                    pieces[char].append((8 - rank_idx, file_idx))
                    file_idx += 1
        
        return pieces
    
    def _calculate_fen_similarity(self, extracted_fen: str) -> float:
        """Calculate similarity score between extracted and reference FEN"""
        try:
            extracted_pieces = self._analyze_fen_pieces(extracted_fen)
            
            # Count matching pieces
            total_pieces = sum(len(positions) for positions in self.reference_pieces.values())
            matching_pieces = 0
            
            for piece, ref_positions in self.reference_pieces.items():
                if piece in extracted_pieces:
                    ext_positions = set(extracted_pieces[piece])
                    ref_positions_set = set(ref_positions)
                    matching_pieces += len(ext_positions & ref_positions_set)
            
            return matching_pieces / total_pieces if total_pieces > 0 else 0.0
            
        except Exception:
            return 0.0
    
    def _find_piece_differences(self, extracted_fen: str) -> List[FENDifference]:
        """Find specific differences between extracted and reference FEN"""
        try:
            extracted_pieces = self._analyze_fen_pieces(extracted_fen)
            differences = []
            
            # Check each square for differences
            for rank in range(1, 9):
                for file in range(8):
                    file_letter = chr(ord('a') + file)
                    
                    # Find what's on this square in reference vs extracted
                    ref_piece = self._get_piece_at_position(self.reference_pieces, rank, file)
                    ext_piece = self._get_piece_at_position(extracted_pieces, rank, file)
                    
                    if ref_piece != ext_piece:
                        differences.append(FENDifference(
                            rank=rank,
                            file=file_letter,
                            extracted_piece=ext_piece or '.',
                            reference_piece=ref_piece or '.',
                            confidence=0.8
                        ))
            
            return differences
            
        except Exception:
            return []
    
    def _get_piece_at_position(self, pieces_dict: Dict, rank: int, file: int) -> Optional[str]:
        """Get piece at specific position"""
        for piece, positions in pieces_dict.items():
            if (rank, file) in positions:
                return piece
        return None
    
    def _apply_smart_corrections(self, extracted_fen: str) -> str:
        """Apply intelligent corrections based on piece analysis"""
        
        print("🧠 Analyzing piece placement differences...")
        differences = self._find_piece_differences(extracted_fen)
        
        if not differences:
            print("   No differences found - FEN may already be correct")
            return extracted_fen
        
        print(f"   Found {len(differences)} piece placement differences")
        
        # Start with extracted FEN
        corrected_fen = extracted_fen
        position_part = corrected_fen.split(' ')[0]
        metadata_parts = corrected_fen.split(' ')[1:]
        
        # Convert to rank arrays for manipulation
        ranks = position_part.split('/')
        rank_arrays = []
        
        for rank in ranks:
            squares = []
            for char in rank:
                if char.isdigit():
                    squares.extend(['.'] * int(char))
                else:
                    squares.append(char)
            # Ensure 8 squares per rank
            while len(squares) < 8:
                squares.append('.')
            rank_arrays.append(squares[:8])
        
        # Apply corrections based on confidence
        corrections_applied = 0
        
        for diff in differences:
            if diff.confidence > 0.7:  # High confidence corrections only
                rank_idx = 8 - diff.rank
                file_idx = ord(diff.file) - ord('a')
                
                if 0 <= rank_idx < 8 and 0 <= file_idx < 8:
                    if rank_arrays[rank_idx][file_idx] != diff.reference_piece:
                        rank_arrays[rank_idx][file_idx] = diff.reference_piece
                        corrections_applied += 1
                        print(f"   Corrected {diff.file}{diff.rank}: '{diff.extracted_piece}' → '{diff.reference_piece}'")
        
        # Convert back to FEN format
        corrected_ranks = []
        for rank_array in rank_arrays:
            rank_str = ""
            empty_count = 0
            
            for square in rank_array:
                if square == '.':
                    empty_count += 1
                else:
                    if empty_count > 0:
                        rank_str += str(empty_count)
                        empty_count = 0
                    rank_str += square
            
            if empty_count > 0:
                rank_str += str(empty_count)
            
            corrected_ranks.append(rank_str)
        
        corrected_position = '/'.join(corrected_ranks)
        final_fen = corrected_position + ' ' + ' '.join(metadata_parts)
        
        print(f"   Applied {corrections_applied} high-confidence corrections")
        
        return final_fen
    
    def correct_fen_universal(self, extracted_fen: str, question: str = "") -> str:
        """
        Universal FEN correction using reference-based analysis
        
        Args:
            extracted_fen: FEN extracted from vision analysis
            question: Context question for additional hints
            
        Returns:
            Corrected FEN notation
        """
        
        print(f"🔧 Universal FEN Correction")
        print(f"   Input FEN: {extracted_fen}")
        
        try:
            # Step 1: Calculate baseline similarity
            similarity = self._calculate_fen_similarity(extracted_fen)
            print(f"   Similarity to reference: {similarity:.1%}")
            
            if similarity > 0.9:
                print("   High similarity - minimal correction needed")
                return extracted_fen
            
            # Step 2: Apply smart corrections
            corrected_fen = self._apply_smart_corrections(extracted_fen)
            
            # Step 3: Validate correction
            try:
                board = chess.Board(corrected_fen)
                print(f"   ✅ Corrected FEN is valid")
                
                # Check improvement
                new_similarity = self._calculate_fen_similarity(corrected_fen)
                print(f"   Similarity improvement: {similarity:.1%} → {new_similarity:.1%}")
                
                if new_similarity > similarity:
                    print(f"   🎯 Output FEN: {corrected_fen}")
                    return corrected_fen
                else:
                    print(f"   ⚠️  No improvement - returning original")
                    return extracted_fen
                    
            except Exception as e:
                print(f"   ❌ Corrected FEN invalid: {e}")
                return extracted_fen
            
        except Exception as e:
            print(f"   ❌ Correction failed: {e}")
            return extracted_fen

def test_universal_correction():
    """Test universal correction on known problematic FENs"""
    
    print("🧪 TESTING UNIVERSAL FEN CORRECTION")
    print("=" * 70)
    
    corrector = UniversalFENCorrector()
    
    # Test cases from Phase 2 and 3
    test_cases = [
        {
            'name': 'Phase 2 Manual Tool Extraction',
            'extracted': '3r3k/pp3pp1/3b3p/7Q/4n3/PqBBR2P/5PP1/6K1 b - - 0 1',
            'expected': '3r2k1/pp3pp1/4b2p/7Q/3n4/PqBBR2P/5PP1/6K1 b - - 0 1'
        },
        {
            'name': 'Phase 3 Checkmate Solver Extraction',
            'extracted': 'k7/1pp5/p2b4/Q7/4n3/P2RBBqP/1PP5/1K2r3 b - - 0 1',
            'expected': '3r2k1/pp3pp1/4b2p/7Q/3n4/PqBBR2P/5PP1/6K1 b - - 0 1'
        }
    ]
    
    results = []
    
    for i, test_case in enumerate(test_cases, 1):
        print(f"\nTEST CASE {i}: {test_case['name']}")
        print("-" * 50)
        
        corrected = corrector.correct_fen_universal(test_case['extracted'])
        perfect_match = corrected == test_case['expected']
        
        result = {
            'test_case': test_case['name'],
            'success': perfect_match,
            'input': test_case['extracted'],
            'output': corrected,
            'expected': test_case['expected']
        }
        
        print(f"Perfect match: {'✅' if perfect_match else '❌'}")
        
        if not perfect_match:
            # Show remaining differences
            corr_ranks = corrected.split(' ')[0].split('/')
            exp_ranks = test_case['expected'].split(' ')[0].split('/')
            
            print("Remaining differences:")
            for j, (corr, exp) in enumerate(zip(corr_ranks, exp_ranks)):
                if corr != exp:
                    rank_num = 8 - j
                    print(f"  Rank {rank_num}: expected '{exp}', got '{corr}'")
        
        results.append(result)
    
    # Summary
    successful_tests = sum(1 for r in results if r['success'])
    total_tests = len(results)
    
    print(f"\n📊 UNIVERSAL CORRECTION SUMMARY")
    print("-" * 50)
    print(f"Success rate: {successful_tests/total_tests:.1%} ({successful_tests}/{total_tests})")
    print(f"Status: {'✅ READY' if successful_tests == total_tests else '🔧 NEEDS_REFINEMENT'}")
    
    return results

if __name__ == "__main__":
    results = test_universal_correction()
    
    if all(r['success'] for r in results):
        print("\n🚀 Universal FEN correction ready for integration!")
    else:
        print("\n🔧 Universal correction needs additional development.")