"""
Phase 1 Auditor - Auditing and validation tools for Phase 2
Verifies labels, detects data leakage, and validates directional accuracy
"""

import pandas as pd
import numpy as np
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple, Any
from datetime import datetime
from loguru import logger
import json


@dataclass
class LabelAuditResult:
    """Result of label verification"""
    horizon: str
    total_samples: int
    valid_samples: int
    invalid_samples: int
    includes_current_bar: bool
    first_invalid_index: Optional[int] = None
    error_rate: float = 0.0
    sample_errors: List[Dict] = field(default_factory=list)


@dataclass
class DirectionalAccuracyResult:
    """Result of directional accuracy calculation"""
    horizon: str
    target_type: str  # 'high' or 'low'
    total_samples: int
    correct_predictions: int
    accuracy: float
    accuracy_by_direction: Dict[str, float] = field(default_factory=dict)


@dataclass
class LeakageCheckResult:
    """Result of data leakage check"""
    check_name: str
    passed: bool
    details: str
    severity: str  # 'critical', 'warning', 'info'
    affected_features: List[str] = field(default_factory=list)


@dataclass
class AuditReport:
    """Complete audit report for Phase 1"""
    timestamp: datetime
    symbol: str
    total_records: int

    # Label verification
    label_results: List[LabelAuditResult] = field(default_factory=list)

    # Directional accuracy
    accuracy_results: List[DirectionalAccuracyResult] = field(default_factory=list)

    # Leakage checks
    leakage_results: List[LeakageCheckResult] = field(default_factory=list)

    # Overall status
    overall_passed: bool = False
    critical_issues: List[str] = field(default_factory=list)
    warnings: List[str] = field(default_factory=list)
    recommendations: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict:
        """Convert report to dictionary"""
        return {
            'timestamp': self.timestamp.isoformat(),
            'symbol': self.symbol,
            'total_records': self.total_records,
            'label_results': [
                {
                    'horizon': r.horizon,
                    'total_samples': r.total_samples,
                    'valid_samples': r.valid_samples,
                    'invalid_samples': r.invalid_samples,
                    'includes_current_bar': r.includes_current_bar,
                    'error_rate': r.error_rate
                }
                for r in self.label_results
            ],
            'accuracy_results': [
                {
                    'horizon': r.horizon,
                    'target_type': r.target_type,
                    'accuracy': r.accuracy,
                    'accuracy_by_direction': r.accuracy_by_direction
                }
                for r in self.accuracy_results
            ],
            'leakage_results': [
                {
                    'check_name': r.check_name,
                    'passed': r.passed,
                    'details': r.details,
                    'severity': r.severity
                }
                for r in self.leakage_results
            ],
            'overall_passed': self.overall_passed,
            'critical_issues': self.critical_issues,
            'warnings': self.warnings,
            'recommendations': self.recommendations
        }

    def to_json(self, filepath: Optional[str] = None) -> str:
        """Export report to JSON"""
        json_str = json.dumps(self.to_dict(), indent=2)
        if filepath:
            with open(filepath, 'w') as f:
                f.write(json_str)
        return json_str

    def print_summary(self):
        """Print human-readable summary"""
        print("\n" + "="*60)
        print("PHASE 1 AUDIT REPORT")
        print("="*60)
        print(f"Symbol: {self.symbol}")
        print(f"Timestamp: {self.timestamp}")
        print(f"Total Records: {self.total_records:,}")
        print(f"Overall Status: {'PASSED' if self.overall_passed else 'FAILED'}")

        print("\n--- Label Verification ---")
        for r in self.label_results:
            status = "OK" if not r.includes_current_bar and r.error_rate == 0 else "ISSUE"
            print(f"  {r.horizon}: {status} (error rate: {r.error_rate:.2%})")

        print("\n--- Directional Accuracy ---")
        for r in self.accuracy_results:
            print(f"  {r.horizon} {r.target_type}: {r.accuracy:.2%}")

        print("\n--- Leakage Checks ---")
        for r in self.leakage_results:
            status = "PASS" if r.passed else "FAIL"
            print(f"  [{r.severity.upper()}] {r.check_name}: {status}")

        if self.critical_issues:
            print("\n--- Critical Issues ---")
            for issue in self.critical_issues:
                print(f"  - {issue}")

        if self.warnings:
            print("\n--- Warnings ---")
            for warning in self.warnings:
                print(f"  - {warning}")

        if self.recommendations:
            print("\n--- Recommendations ---")
            for rec in self.recommendations:
                print(f"  - {rec}")

        print("="*60 + "\n")


class Phase1Auditor:
    """
    Auditor for Phase 1 models and data pipeline

    Performs:
    1. Label verification (future High/Low calculation)
    2. Directional accuracy recalculation
    3. Data leakage detection
    """

    # Horizon configurations for Phase 2
    HORIZONS = {
        '15m': {'bars': 3, 'start': 1, 'end': 3},
        '1h': {'bars': 12, 'start': 1, 'end': 12}
    }

    def __init__(self):
        """Initialize auditor"""
        self.report = None

    def run_full_audit(
        self,
        df: pd.DataFrame,
        symbol: str,
        predictions: Optional[pd.DataFrame] = None
    ) -> AuditReport:
        """
        Run complete audit on data and predictions

        Args:
            df: DataFrame with OHLCV data
            symbol: Trading symbol
            predictions: Optional DataFrame with model predictions

        Returns:
            AuditReport with all findings
        """
        logger.info(f"Starting full audit for {symbol}")

        self.report = AuditReport(
            timestamp=datetime.now(),
            symbol=symbol,
            total_records=len(df)
        )

        # 1. Verify labels
        self._verify_labels(df)

        # 2. Check directional accuracy (if predictions provided)
        if predictions is not None:
            self._check_directional_accuracy(df, predictions)

        # 3. Detect data leakage
        self._detect_data_leakage(df)

        # 4. Generate recommendations
        self._generate_recommendations()

        # 5. Determine overall status
        self.report.overall_passed = (
            len(self.report.critical_issues) == 0 and
            all(r.passed for r in self.report.leakage_results if r.severity == 'critical')
        )

        logger.info(f"Audit completed. Status: {'PASSED' if self.report.overall_passed else 'FAILED'}")
        return self.report

    def verify_future_labels(
        self,
        df: pd.DataFrame,
        horizon_name: str = '15m'
    ) -> LabelAuditResult:
        """
        Verify that future labels are calculated correctly

        Labels should be:
        - high_15m = max(high[t+1 ... t+3])  # NOT including t
        - low_15m  = min(low[t+1 ... t+3])
        - high_1h  = max(high[t+1 ... t+12])
        - low_1h   = min(low[t+1 ... t+12])

        Args:
            df: DataFrame with OHLCV data
            horizon_name: Horizon to verify ('15m' or '1h')

        Returns:
            LabelAuditResult with verification details
        """
        config = self.HORIZONS[horizon_name]
        start_offset = config['start']
        end_offset = config['end']

        logger.info(f"Verifying labels for {horizon_name} (bars {start_offset} to {end_offset})")

        # Calculate correct labels
        correct_high = self._calculate_future_max(df['high'], start_offset, end_offset)
        correct_low = self._calculate_future_min(df['low'], start_offset, end_offset)

        # Check if existing labels include current bar (t=0)
        # This would be wrong: max(high[t ... t+3]) instead of max(high[t+1 ... t+3])
        wrong_high = self._calculate_future_max(df['high'], 0, end_offset)
        wrong_low = self._calculate_future_min(df['low'], 0, end_offset)

        # Check for existing label columns
        high_col = f'future_high_{horizon_name}'
        low_col = f'future_low_{horizon_name}'

        includes_current = False
        invalid_samples = 0
        sample_errors = []

        if high_col in df.columns:
            # Check if labels match correct calculation
            mask_valid = ~df[high_col].isna() & ~correct_high.isna()

            # Check if they match wrong calculation (including current bar)
            matches_wrong = np.allclose(
                df.loc[mask_valid, high_col].values,
                wrong_high.loc[mask_valid].values,
                rtol=1e-5, equal_nan=True
            )

            matches_correct = np.allclose(
                df.loc[mask_valid, high_col].values,
                correct_high.loc[mask_valid].values,
                rtol=1e-5, equal_nan=True
            )

            if matches_wrong and not matches_correct:
                includes_current = True
                invalid_samples = mask_valid.sum()
                logger.warning(f"Labels for {horizon_name} include current bar (t=0)!")
            elif not matches_correct:
                # Find mismatches
                diff = abs(df.loc[mask_valid, high_col] - correct_high.loc[mask_valid])
                mismatches = diff > 1e-5
                invalid_samples = mismatches.sum()

                # Sample some errors
                if invalid_samples > 0:
                    error_indices = diff[mismatches].nsmallest(5).index.tolist()
                    for idx in error_indices:
                        sample_errors.append({
                            'index': str(idx),
                            'existing': float(df.loc[idx, high_col]),
                            'correct': float(correct_high.loc[idx]),
                            'diff': float(diff.loc[idx])
                        })

        result = LabelAuditResult(
            horizon=horizon_name,
            total_samples=len(df),
            valid_samples=len(df) - invalid_samples,
            invalid_samples=invalid_samples,
            includes_current_bar=includes_current,
            error_rate=invalid_samples / len(df) if len(df) > 0 else 0,
            sample_errors=sample_errors
        )

        return result

    def calculate_correct_labels(
        self,
        df: pd.DataFrame,
        horizon_name: str = '15m'
    ) -> pd.DataFrame:
        """
        Calculate correct future labels (not including current bar)

        Args:
            df: DataFrame with OHLCV data
            horizon_name: Horizon name ('15m' or '1h')

        Returns:
            DataFrame with correct labels added
        """
        df = df.copy()
        config = self.HORIZONS[horizon_name]
        start_offset = config['start']
        end_offset = config['end']

        # Calculate correct labels (starting from t+1, NOT t)
        df[f'future_high_{horizon_name}'] = self._calculate_future_max(
            df['high'], start_offset, end_offset
        )
        df[f'future_low_{horizon_name}'] = self._calculate_future_min(
            df['low'], start_offset, end_offset
        )

        # Calculate delta (range) targets for Phase 2
        df[f'delta_high_{horizon_name}'] = df[f'future_high_{horizon_name}'] - df['close']
        df[f'delta_low_{horizon_name}'] = df['close'] - df[f'future_low_{horizon_name}']

        logger.info(f"Calculated correct labels for {horizon_name}")
        return df

    def check_directional_accuracy(
        self,
        df: pd.DataFrame,
        predictions: pd.DataFrame,
        horizon_name: str = '15m'
    ) -> Tuple[DirectionalAccuracyResult, DirectionalAccuracyResult]:
        """
        Calculate directional accuracy correctly

        For High predictions:
            sign(pred_high - close_t) == sign(real_high - close_t)

        For Low predictions:
            sign(close_t - pred_low) == sign(close_t - real_low)

        Args:
            df: DataFrame with OHLCV and actual future values
            predictions: DataFrame with predicted values
            horizon_name: Horizon name

        Returns:
            Tuple of (high_accuracy_result, low_accuracy_result)
        """
        # Get actual and predicted values
        actual_high = df[f'future_high_{horizon_name}']
        actual_low = df[f'future_low_{horizon_name}']
        close = df['close']

        pred_high_col = f'pred_high_{horizon_name}'
        pred_low_col = f'pred_low_{horizon_name}'

        # Check if prediction columns exist
        if pred_high_col not in predictions.columns or pred_low_col not in predictions.columns:
            logger.warning(f"Prediction columns not found for {horizon_name}")
            return None, None

        pred_high = predictions[pred_high_col]
        pred_low = predictions[pred_low_col]

        # Align indices
        common_idx = df.index.intersection(predictions.index)

        # High directional accuracy
        # sign(pred_high - close_t) == sign(real_high - close_t)
        sign_pred_high = np.sign(pred_high.loc[common_idx] - close.loc[common_idx])
        sign_real_high = np.sign(actual_high.loc[common_idx] - close.loc[common_idx])

        high_correct = (sign_pred_high == sign_real_high)
        high_accuracy = high_correct.mean()

        # Accuracy by direction
        high_acc_up = high_correct[sign_real_high > 0].mean() if (sign_real_high > 0).any() else 0
        high_acc_down = high_correct[sign_real_high < 0].mean() if (sign_real_high < 0).any() else 0

        high_result = DirectionalAccuracyResult(
            horizon=horizon_name,
            target_type='high',
            total_samples=len(common_idx),
            correct_predictions=high_correct.sum(),
            accuracy=high_accuracy,
            accuracy_by_direction={'up': high_acc_up, 'down': high_acc_down}
        )

        # Low directional accuracy
        # sign(close_t - pred_low) == sign(close_t - real_low)
        sign_pred_low = np.sign(close.loc[common_idx] - pred_low.loc[common_idx])
        sign_real_low = np.sign(close.loc[common_idx] - actual_low.loc[common_idx])

        low_correct = (sign_pred_low == sign_real_low)
        low_accuracy = low_correct.mean()

        # Accuracy by direction
        low_acc_up = low_correct[sign_real_low > 0].mean() if (sign_real_low > 0).any() else 0
        low_acc_down = low_correct[sign_real_low < 0].mean() if (sign_real_low < 0).any() else 0

        low_result = DirectionalAccuracyResult(
            horizon=horizon_name,
            target_type='low',
            total_samples=len(common_idx),
            correct_predictions=low_correct.sum(),
            accuracy=low_accuracy,
            accuracy_by_direction={'up': low_acc_up, 'down': low_acc_down}
        )

        return high_result, low_result

    def detect_data_leakage(self, df: pd.DataFrame) -> List[LeakageCheckResult]:
        """
        Detect potential data leakage issues

        Checks:
        1. Temporal ordering
        2. Centered rolling windows
        3. Future-looking features

        Args:
            df: DataFrame to check

        Returns:
            List of LeakageCheckResult
        """
        results = []

        # Check 1: Temporal ordering
        if df.index.is_monotonic_increasing:
            results.append(LeakageCheckResult(
                check_name="Temporal Ordering",
                passed=True,
                details="Index is monotonically increasing (correct)",
                severity="critical"
            ))
        else:
            results.append(LeakageCheckResult(
                check_name="Temporal Ordering",
                passed=False,
                details="Index is NOT monotonically increasing - data may be shuffled!",
                severity="critical"
            ))

        # Check 2: Look for centered rolling calculations
        # These would have NaN at both ends instead of just the beginning
        for col in df.columns:
            if 'roll' in col.lower() or 'ma' in col.lower() or 'avg' in col.lower():
                nan_start = df[col].isna().iloc[:50].sum()
                nan_end = df[col].isna().iloc[-50:].sum()

                if nan_end > nan_start:
                    results.append(LeakageCheckResult(
                        check_name=f"Centered Window: {col}",
                        passed=False,
                        details=f"Column {col} may use centered window (NaN at end: {nan_end})",
                        severity="critical",
                        affected_features=[col]
                    ))

        # Check 3: Look for future-looking column names
        future_keywords = ['future', 'next', 'forward', 'target', 'label']
        feature_cols = [c for c in df.columns if not any(kw in c.lower() for kw in ['t_', 'future_'])]

        suspicious_features = []
        for col in feature_cols:
            for kw in future_keywords:
                if kw in col.lower():
                    suspicious_features.append(col)

        if suspicious_features:
            results.append(LeakageCheckResult(
                check_name="Future-Looking Features",
                passed=False,
                details=f"Found potentially future-looking features in non-target columns",
                severity="warning",
                affected_features=suspicious_features
            ))
        else:
            results.append(LeakageCheckResult(
                check_name="Future-Looking Features",
                passed=True,
                details="No suspicious future-looking features found",
                severity="info"
            ))

        # Check 4: Duplicate timestamps
        if df.index.duplicated().any():
            n_dups = df.index.duplicated().sum()
            results.append(LeakageCheckResult(
                check_name="Duplicate Timestamps",
                passed=False,
                details=f"Found {n_dups} duplicate timestamps",
                severity="warning"
            ))
        else:
            results.append(LeakageCheckResult(
                check_name="Duplicate Timestamps",
                passed=True,
                details="No duplicate timestamps found",
                severity="info"
            ))

        return results

    def validate_scaler_usage(
        self,
        train_data: pd.DataFrame,
        val_data: pd.DataFrame,
        scaler_fit_data: pd.DataFrame
    ) -> LeakageCheckResult:
        """
        Validate that scaler was fit only on training data

        Args:
            train_data: Training data
            val_data: Validation data
            scaler_fit_data: Data that scaler was fitted on

        Returns:
            LeakageCheckResult
        """
        # Check if scaler_fit_data matches train_data
        if len(scaler_fit_data) > len(train_data):
            return LeakageCheckResult(
                check_name="Scaler Fit Data",
                passed=False,
                details="Scaler was fit on more data than training set - possible leakage!",
                severity="critical"
            )

        # Check if validation data indices are in fit data
        common_idx = val_data.index.intersection(scaler_fit_data.index)
        if len(common_idx) > 0:
            return LeakageCheckResult(
                check_name="Scaler Fit Data",
                passed=False,
                details=f"Scaler fit data contains {len(common_idx)} validation samples!",
                severity="critical"
            )

        return LeakageCheckResult(
            check_name="Scaler Fit Data",
            passed=True,
            details="Scaler was correctly fit only on training data",
            severity="critical"
        )

    def validate_walk_forward_split(
        self,
        train_indices: np.ndarray,
        val_indices: np.ndarray,
        test_indices: np.ndarray
    ) -> LeakageCheckResult:
        """
        Validate that walk-forward split is strictly temporal

        Args:
            train_indices: Training set indices (as timestamps or integers)
            val_indices: Validation set indices
            test_indices: Test set indices

        Returns:
            LeakageCheckResult
        """
        # Check train < val < test
        train_max = np.max(train_indices)
        val_min = np.min(val_indices)
        val_max = np.max(val_indices)
        test_min = np.min(test_indices)

        issues = []

        if train_max >= val_min:
            issues.append(f"Train max ({train_max}) >= Val min ({val_min})")

        if val_max >= test_min:
            issues.append(f"Val max ({val_max}) >= Test min ({test_min})")

        # Check for overlaps
        train_val_overlap = np.intersect1d(train_indices, val_indices)
        val_test_overlap = np.intersect1d(val_indices, test_indices)
        train_test_overlap = np.intersect1d(train_indices, test_indices)

        if len(train_val_overlap) > 0:
            issues.append(f"Train-Val overlap: {len(train_val_overlap)} samples")

        if len(val_test_overlap) > 0:
            issues.append(f"Val-Test overlap: {len(val_test_overlap)} samples")

        if len(train_test_overlap) > 0:
            issues.append(f"Train-Test overlap: {len(train_test_overlap)} samples")

        if issues:
            return LeakageCheckResult(
                check_name="Walk-Forward Split",
                passed=False,
                details="; ".join(issues),
                severity="critical"
            )

        return LeakageCheckResult(
            check_name="Walk-Forward Split",
            passed=True,
            details="Walk-forward split is strictly temporal with no overlaps",
            severity="critical"
        )

    # Private helper methods

    def _calculate_future_max(
        self,
        series: pd.Series,
        start_offset: int,
        end_offset: int
    ) -> pd.Series:
        """Calculate max of future values (not including current)"""
        future_values = []
        for i in range(start_offset, end_offset + 1):
            future_values.append(series.shift(-i))
        return pd.concat(future_values, axis=1).max(axis=1)

    def _calculate_future_min(
        self,
        series: pd.Series,
        start_offset: int,
        end_offset: int
    ) -> pd.Series:
        """Calculate min of future values (not including current)"""
        future_values = []
        for i in range(start_offset, end_offset + 1):
            future_values.append(series.shift(-i))
        return pd.concat(future_values, axis=1).min(axis=1)

    def _verify_labels(self, df: pd.DataFrame):
        """Verify labels for all horizons"""
        for horizon_name in self.HORIZONS.keys():
            result = self.verify_future_labels(df, horizon_name)
            self.report.label_results.append(result)

            if result.includes_current_bar:
                self.report.critical_issues.append(
                    f"Labels for {horizon_name} include current bar (t=0)"
                )

    def _check_directional_accuracy(self, df: pd.DataFrame, predictions: pd.DataFrame):
        """Check directional accuracy for all horizons"""
        for horizon_name in self.HORIZONS.keys():
            high_result, low_result = self.check_directional_accuracy(
                df, predictions, horizon_name
            )
            if high_result:
                self.report.accuracy_results.append(high_result)
            if low_result:
                self.report.accuracy_results.append(low_result)

    def _detect_data_leakage(self, df: pd.DataFrame):
        """Run all leakage detection checks"""
        leakage_results = self.detect_data_leakage(df)
        self.report.leakage_results.extend(leakage_results)

        for result in leakage_results:
            if not result.passed:
                if result.severity == 'critical':
                    self.report.critical_issues.append(
                        f"[{result.check_name}] {result.details}"
                    )
                elif result.severity == 'warning':
                    self.report.warnings.append(
                        f"[{result.check_name}] {result.details}"
                    )

    def _generate_recommendations(self):
        """Generate recommendations based on findings"""
        # Based on label issues
        for result in self.report.label_results:
            if result.includes_current_bar:
                self.report.recommendations.append(
                    f"Recalculate {result.horizon} labels to exclude current bar (use t+1 to t+n)"
                )
            elif result.error_rate > 0:
                self.report.recommendations.append(
                    f"Review {result.horizon} label calculation - {result.error_rate:.2%} error rate"
                )

        # Based on accuracy imbalance
        for result in self.report.accuracy_results:
            if result.target_type == 'high' and result.accuracy > 0.9:
                self.report.recommendations.append(
                    f"High accuracy for {result.horizon} high predictions ({result.accuracy:.2%}) "
                    "may indicate data leakage - verify calculation"
                )
            elif result.target_type == 'low' and result.accuracy < 0.2:
                self.report.recommendations.append(
                    f"Low accuracy for {result.horizon} low predictions ({result.accuracy:.2%}) - "
                    "verify directional accuracy formula"
                )

        # Based on leakage
        for result in self.report.leakage_results:
            if not result.passed and result.affected_features:
                self.report.recommendations.append(
                    f"Review features: {', '.join(result.affected_features)}"
                )


if __name__ == "__main__":
    # Test the auditor
    import numpy as np

    # Create sample data
    np.random.seed(42)
    n_samples = 1000

    dates = pd.date_range(start='2023-01-01', periods=n_samples, freq='5min')

    df = pd.DataFrame({
        'open': np.random.randn(n_samples).cumsum() + 100,
        'high': np.random.randn(n_samples).cumsum() + 101,
        'low': np.random.randn(n_samples).cumsum() + 99,
        'close': np.random.randn(n_samples).cumsum() + 100,
        'volume': np.random.randint(1000, 10000, n_samples)
    }, index=dates)

    # Make high/low consistent
    df['high'] = df[['open', 'close']].max(axis=1) + abs(np.random.randn(n_samples) * 0.5)
    df['low'] = df[['open', 'close']].min(axis=1) - abs(np.random.randn(n_samples) * 0.5)

    # Run audit
    auditor = Phase1Auditor()
    report = auditor.run_full_audit(df, symbol='TEST')

    # Print summary
    report.print_summary()

    # Test label calculation
    df_with_labels = auditor.calculate_correct_labels(df, '15m')
    print("\nSample labels:")
    print(df_with_labels[['close', 'future_high_15m', 'future_low_15m',
                          'delta_high_15m', 'delta_low_15m']].head(10))
