scripts/data_quality_validator.py

#!/usr/bin/env python3
"""
Data Quality Validator for Stock Predictor
Validates data quality across all data sources and features

Usage:
    python data_quality_validator.py [--check all|features|stock_data|predictions]
    python data_quality_validator.py --report
"""

import os
import sys
import json
from pathlib import Path
from datetime import datetime
from typing import Dict, List, Any, Optional
from dataclasses import dataclass, field, asdict
import pandas as pd
import numpy as np

# Add project root to path
PROJECT_DIR = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_DIR))


@dataclass
class ValidationResult:
    """Data quality validation result"""
    check_name: str
    passed: bool
    message: str
    details: List[str] = field(default_factory=list)
    severity: str = "info"  # info, warning, error
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())


class DataQualityValidator:
    """Data quality validator for stock predictor"""

    def __init__(self, project_dir: Path = PROJECT_DIR):
        self.project_dir = project_dir
        self.data_dir = project_dir / "data"
        self.results: List[ValidationResult] = []

    def add_result(self, result: ValidationResult):
        """Add validation result"""
        self.results.append(result)
        status = "✓" if result.passed else "✗"
        severity = result.severity.upper()
        print(f"{status} [{severity}] {result.check_name}: {result.message}")
        if result.details:
            for detail in result.details[:3]:  # Show first 3 details
                print(f"    - {detail}")

    def validate_json_structure(self, filepath: Path, required_keys: List[str]) -> ValidationResult:
        """Validate JSON file structure"""
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                data = json.load(f)

            missing_keys = [k for k in required_keys if k not in data]
            if missing_keys:
                return ValidationResult(
                    check_name=f"JSON structure: {filepath.name}",
                    passed=False,
                    message=f"Missing required keys: {missing_keys}",
                    details=[f"File: {filepath}"],
                    severity="error"
                )

            return ValidationResult(
                check_name=f"JSON structure: {filepath.name}",
                passed=True,
                message="All required keys present",
                details=[f"File: {filepath}, Keys: {len(data.keys())}"],
                severity="info"
            )

        except json.JSONDecodeError as e:
            return ValidationResult(
                check_name=f"JSON structure: {filepath.name}",
                passed=False,
                message=f"Invalid JSON: {e}",
                details=[f"File: {filepath}"],
                severity="error"
            )

        except Exception as e:
            return ValidationResult(
                check_name=f"JSON structure: {filepath.name}",
                passed=False,
                message=f"Cannot read file: {e}",
                details=[f"File: {filepath}"],
                severity="error"
            )

    def validate_stock_data(self) -> List[ValidationResult]:
        """Validate stock data files"""
        results = []

        # Check TAIEX data
        taifex_path = self.data_dir / "taifex" / "daily"
        if taifex_path.exists():
            files = list(taifex_path.glob("*.csv"))
            if files:
                latest_file = max(files, key=lambda f: f.stat().st_mtime)
                try:
                    df = pd.read_csv(latest_file)
                    if df.shape[0] > 0:
                        results.append(ValidationResult(
                            check_name="TAIEX daily data",
                            passed=True,
                            message=f"Loaded {df.shape[0]} rows, {df.shape[1]} columns",
                            details=[f"Latest file: {latest_file.name}"],
                            severity="info"
                        ))
                    else:
                        results.append(ValidationResult(
                            check_name="TAIEX daily data",
                            passed=False,
                            message="Empty dataframe",
                            details=[f"File: {latest_file.name}"],
                            severity="warning"
                        ))
                except Exception as e:
                    results.append(ValidationResult(
                        check_name="TAIEX daily data",
                        passed=False,
                        message=f"Cannot read CSV: {e}",
                        details=[f"File: {latest_file.name}"],
                        severity="error"
                    ))
            else:
                results.append(ValidationResult(
                    check_name="TAIEX daily data",
                    passed=False,
                    message="No CSV files found",
                    details=[f"Directory: {taifex_path}"],
                    severity="warning"
                ))

        return results

    def validate_predictions(self) -> List[ValidationResult]:
        """Validate prediction files"""
        results = []

        # Check precomputed predictions
        predictions_path = self.data_dir / "precomputed_predictions.json"
        if predictions_path.exists():
            result = self.validate_json_structure(predictions_path, [])
            result.details.append(f"Size: {predictions_path.stat().st_size / 1024:.1f} KB")
            results.append(result)
        else:
            results.append(ValidationResult(
                check_name="Precomputed predictions",
                passed=False,
                message="File not found",
                details=[f"Path: {predictions_path}"],
                severity="warning"
            ))

        return results

    def validate_query_history(self) -> List[ValidationResult]:
        """Validate query history file"""
        results = []

        history_path = self.data_dir / "query_history.json"
        if history_path.exists():
            result = self.validate_json_structure(history_path, ["predictions"])
            if result.passed:
                with open(history_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    if "predictions" in data:
                        results.append(ValidationResult(
                            check_name="Query history predictions",
                            passed=True,
                            message=f"Contains {len(data['predictions'])} predictions",
                            details=[f"File: {history_path}"],
                            severity="info"
                        ))

        return results

    def validate_feature_columns(self) -> List[ValidationResult]:
        """Validate feature columns in models"""
        results = []

        try:
            # Import the predictor module
            from models.predictor import FEATURE_COLUMNS, FEATURE_VERSION

            # Check if features are defined
            if FEATURE_COLUMNS:
                results.append(ValidationResult(
                    check_name="Feature columns defined",
                    passed=True,
                    message=f"{len(FEATURE_COLUMNS)} features defined (v{FEATURE_VERSION})",
                    details=[f"Feature count: {len(FEATURE_COLUMNS)}", f"Version: {FEATURE_VERSION}"],
                    severity="info"
                ))
            else:
                results.append(ValidationResult(
                    check_name="Feature columns defined",
                    passed=False,
                    message="No features defined",
                    details=[],
                    severity="error"
                ))

        except Exception as e:
            results.append(ValidationResult(
                check_name="Feature columns defined",
                passed=False,
                message=f"Cannot import features: {e}",
                details=[],
                severity="error"
            ))

        return results

    def generate_report(self, output_path: Optional[Path] = None) -> Dict[str, Any]:
        """Generate validation report"""

        # Calculate summary
        total = len(self.results)
        passed = sum(1 for r in self.results if r.passed)
        failed = total - passed

        # Group by severity
        by_severity = {"info": 0, "warning": 0, "error": 0}
        for r in self.results:
            by_severity[r.severity] += 1

        report = {
            "timestamp": datetime.now().isoformat(),
            "summary": {
                "total_checks": total,
                "passed": passed,
                "failed": failed,
                "by_severity": by_severity
            },
            "results": [asdict(r) for r in self.results]
        }

        if output_path:
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(report, f, indent=2, ensure_ascii=False)

        return report


def main():
    """Main function"""
    import argparse

    parser = argparse.ArgumentParser(description="Data Quality Validator")
    parser.add_argument(
        "--check",
        choices=["all", "features", "stock_data", "predictions", "query_history"],
        default="all",
        help="Check type"
    )
    parser.add_argument(
        "--report",
        action="store_true",
        help="Generate detailed report"
    )
    parser.add_argument(
        "--output",
        help="Output path for report"
    )

    args = parser.parse_args()

    validator = DataQualityValidator()

    if args.check == "features":
        results = validator.validate_feature_columns()
    elif args.check == "stock_data":
        results = validator.validate_stock_data()
    elif args.check == "predictions":
        results = validator.validate_predictions()
    elif args.check == "query_history":
        results = validator.validate_query_history()
    else:  # all
        results = (
            validator.validate_feature_columns() +
            validator.validate_stock_data() +
            validator.validate_predictions() +
            validator.validate_query_history()
        )

    for r in results:
        validator.add_result(r)

    # Generate report if requested
    if args.report:
        report = validator.generate_report()
        if args.output:
            output_path = Path(args.output)
            output_path.parent.mkdir(parents=True, exist_ok=True)
        else:
            output_path = PROJECT_DIR / "docs" / "data_quality_report.json"

        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(report, f, indent=2, ensure_ascii=False)

        print(f"\n✓ Report saved to: {output_path}")
        print(f"\nSummary:")
        print(f"  Total checks: {report['summary']['total_checks']}")
        print(f"  Passed: {report['summary']['passed']}")
        print(f"  Failed: {report['summary']['failed']}")
        print(f"  Errors: {report['summary']['by_severity']['error']}")
        print(f"  Warnings: {report['summary']['by_severity']['warning']}")
        print(f"  Info: {report['summary']['by_severity']['info']}")

    # Exit with error code if any errors
    exit_code = 0
    for r in validator.results:
        if r.severity == "error" and not r.passed:
            exit_code = 1
            break

    return exit_code


if __name__ == "__main__":
    exit(main())