"""
Integration and Export Module

Export well log interpretations to industry-standard formats compatible with
Petrel, Techlog, Interactive Petrophysics, and other interpretation software.

Implements bidirectional data exchange for human-in-the-loop workflows.
"""

import datetime
import logging
from pathlib import Path

import numpy as np
import pandas as pd

# FormationTop not found, may need to import from different location
# from dataclasses import dataclass


@dataclass
class FormationTop:
    """Formation top marker."""

    well_name: str
    formation_name: str
    depth: float
    confidence: float | None = None
    notes: str | None = None


# Configure logging
logger = logging.getLogger(__name__)


class LASExporter:
    """
    Export interpreted logs back to LAS format.

    Adds interpreted curves (facies, formation codes) to original log data
    for import into interpretation software.
    """

    def __init__(self, version: str = "2.0"):
        """
        Initialize LAS exporter.

        Args:
            version: LAS version ('2.0' or '3.0')
        """
        self.version = version

    def export_with_interpretation(
        self,
        output_file: str,
        depth: np.ndarray,
        original_curves: pd.DataFrame,
        interpreted_curves: dict[str, np.ndarray],
        well_header: dict[str, str] | None = None,
        curve_descriptions: dict[str, str] | None = None,
        null_value: float = -999.25,
    ):
        """
        Export well logs with interpreted curves to LAS file.

        Args:
            output_file: Path to output LAS file
            depth: Depth array
            original_curves: DataFrame with original log curves
            interpreted_curves: Dict of {curve_name: values} for interpreted data
            well_header: Optional well header information
            curve_descriptions: Optional descriptions for interpreted curves
            null_value: Null value indicator
        """
        with open(output_file, "w") as f:
            # Write version section
            f.write("~VERSION INFORMATION\n")
            f.write(
                f" VERS.                  {self.version}:   "
                f"CWLS LOG ASCII STANDARD - VERSION {self.version}\n"
            )
            f.write(" WRAP.                  NO:   ONE LINE PER DEPTH STEP\n")
            f.write(
                " GENERATED BY.          PyGeomodeling:   "
                "Automated Well Log Interpretation\n"
            )
            f.write(
                f" DATE.                  "
                f"{datetime.datetime.now().strftime('%Y-%m-%d')}:   "
                f"Export Date\n"
            )

            # Write well information section
            f.write("~WELL INFORMATION\n")
            if well_header:
                f.write(
                    f" WELL.                  "
                    f"{well_header.get('well_name', 'UNKNOWN')}:   Well Name\n"
                )
                f.write(
                    f" UWI.                   "
                    f"{well_header.get('uwi', '')}:   Unique Well Identifier\n"
                )
                f.write(
                    f" FLD.                   "
                    f"{well_header.get('field', '')}:   Field\n"
                )
                f.write(
                    f" LOC.                   "
                    f"{well_header.get('location', '')}:   Location\n"
                )
                f.write(
                    f" COMP.                  "
                    f"{well_header.get('operator', '')}:   Operator\n"
                )
            f.write(f" STRT.M                 {depth.min():.2f}:   Start Depth\n")
            f.write(f" STOP.M                 {depth.max():.2f}:   Stop Depth\n")
            f.write(
                f" STEP.M                 {np.median(np.diff(depth)):.4f}:   Step\n"
            )
            f.write(f" NULL.                  {null_value}:   Null Value\n")

            # Write curve information section
            f.write("~CURVE INFORMATION\n")
            f.write(" DEPT.M                                    :  Depth\n")

            # Original curves
            for col in original_curves.columns:
                unit = "unitless"
                desc = col
                f.write(f" {col:6s}.{unit:20s}:  {desc}\n")

            # Interpreted curves
            for curve_name in interpreted_curves.keys():
                unit = "code"
                desc = (
                    curve_descriptions.get(curve_name, curve_name)
                    if curve_descriptions
                    else curve_name
                )
                f.write(f" {curve_name:6s}.{unit:20s}:  {desc}\n")

            # Write parameter section (optional)
            f.write("~PARAMETER INFORMATION\n")

            # Write data section
            f.write("~ASCII LOG DATA\n")

            # Combine all data
            for i, d in enumerate(depth):
                row = [f"{d:10.4f}"]

                # Original curves
                for col in original_curves.columns:
                    val = (
                        original_curves[col].iloc[i]
                        if i < len(original_curves)
                        else null_value
                    )
                    row.append(f"{val:10.4f}")

                # Interpreted curves
                for curve_name in interpreted_curves.keys():
                    val = (
                        interpreted_curves[curve_name][i]
                        if i < len(interpreted_curves[curve_name])
                        else null_value
                    )
                    row.append(f"{val:10.0f}")

                f.write(" ".join(row) + "\n")

        logger.info("Exported interpreted logs to %s", output_file)


class FormationTopExporter:
    """Export formation tops in various formats."""

    @staticmethod
    def export_to_csv(
        tops: list[FormationTop],
        output_file: str,
        well_name: str,
    ):
        """
        Export formation tops to CSV format.

        Args:
            tops: List of FormationTop objects
            output_file: Path to output CSV
            well_name: Well identifier
        """
        rows = []
        for top in tops:
            rows.append(
                {
                    "Well": well_name,
                    "Formation": top.formation_name,
                    "Depth_m": top.depth,
                    "Confidence": top.confidence,
                    "Method": top.method,
                }
            )

        df = pd.DataFrame(rows)
        df.to_csv(output_file, index=False)
        logger.info("Exported %d formation tops to %s", len(tops), output_file)

    @staticmethod
    def export_to_petrel_format(
        tops_by_well: dict[str, list[FormationTop]],
        output_file: str,
    ):
        """
        Export formation tops in Petrel-compatible format.

        Format: Well, X, Y, Z, Formation, MD, TVD

        Args:
            tops_by_well: Dict of {well_name: [FormationTop]}
            output_file: Path to output file
        """
        with open(output_file, "w") as f:
            # Header
            f.write("# Petrel Formation Tops Export\n")
            f.write("# Generated by PyGeomodeling\n")
            f.write("Well\tFormation\tMD\tTVD\tConfidence\tMethod\n")

            for well_name, tops in tops_by_well.items():
                for top in tops:
                    f.write(
                        f"{well_name}\t{top.formation_name}\t{top.depth:.2f}\t{top.depth:.2f}\t{top.confidence:.3f}\t{top.method}\n"
                    )

        logger.info(
            "Exported formation tops for %d wells to %s", len(tops_by_well), output_file
        )

    @staticmethod
    def export_to_ascii(
        tops: list[FormationTop],
        output_file: str,
        well_name: str,
        uwi: str | None = None,
    ):
        """
        Export to simple ASCII format for various software.

        Args:
            tops: List of FormationTop objects
            output_file: Path to output file
            well_name: Well identifier
            uwi: Unique Well Identifier
        """
        with open(output_file, "w") as f:
            f.write(f"WELL: {well_name}\n")
            if uwi:
                f.write(f"UWI: {uwi}\n")
            f.write(f"DATE: {datetime.datetime.now().strftime('%Y-%m-%d')}\n")
            f.write(
                "SOURCE: PyGeomodeling Automated Interpretation\n"
            )
            f.write("\n")
            f.write(
                f"{'Formation':<30} {'Depth (m)':>12} "
                f"{'Confidence':>12} {'Method':<20}\n"
            )
            f.write("-" * 80 + "\n")

            for top in tops:
                f.write(
                    f"{top.formation_name:<30} {top.depth:>12.2f} "
                    f"{top.confidence:>12.3f} {top.method:<20}\n"
                )

        logger.info("Exported formation tops to %s", output_file)


class FaciesLogExporter:
    """Export facies interpretations in various formats."""

    @staticmethod
    def export_to_csv(
        depth: np.ndarray,
        facies: np.ndarray,
        confidence: np.ndarray | None = None,
        output_file: str = "facies_log.csv",
        facies_names: dict[int, str] | None = None,
    ):
        """
        Export facies log to CSV.

        Args:
            depth: Depth array
            facies: Facies codes
            confidence: Optional confidence scores
            output_file: Path to output CSV
            facies_names: Optional mapping of codes to names
        """
        data = {
            "Depth_m": depth,
            "Facies_Code": facies,
        }

        if facies_names:
            data["Facies_Name"] = [facies_names.get(f, f"Unknown_{f}") for f in facies]

        if confidence is not None:
            data["Confidence"] = confidence

        df = pd.DataFrame(data)
        df.to_csv(output_file, index=False)
        logger.info(
            "Exported facies log with %d samples to %s", len(depth), output_file
        )

    @staticmethod
    def create_facies_curve_for_las(
        facies: np.ndarray,
        depth: np.ndarray,
    ) -> dict[str, np.ndarray]:
        """
        Create facies curve dictionary for LAS export.

        Args:
            facies: Facies codes
            depth: Depth array

        Returns:
            Dictionary of curve_name: values
        """
        return {"FACIES": facies.astype(float)}


class PetrelProjectExporter:
    """
    Export complete interpretation package for Petrel import.

    Creates a folder structure with all necessary files.
    """

    @staticmethod
    def export_interpretation_package(
        output_dir: str,
        well_name: str,
        depth: np.ndarray,
        original_curves: pd.DataFrame,
        facies: np.ndarray | None = None,
        formation_tops: list[FormationTop] | None = None,
        confidence_scores: np.ndarray | None = None,
        well_header: dict[str, str] | None = None,
    ):
        """
        Export complete interpretation package.

        Creates directory structure:
        - interpreted_logs/well_name.las
        - formation_tops/well_name_tops.txt
        - qc/confidence_report.csv

        Args:
            output_dir: Base output directory
            well_name: Well identifier
            depth: Depth array
            original_curves: Original log curves
            facies: Interpreted facies (optional)
            formation_tops: Interpreted formation tops (optional)
            confidence_scores: Prediction confidence (optional)
            well_header: Well metadata (optional)
        """
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)

        # Create subdirectories
        (output_path / "interpreted_logs").mkdir(exist_ok=True)
        (output_path / "formation_tops").mkdir(exist_ok=True)
        (output_path / "qc").mkdir(exist_ok=True)

        # Export interpreted LAS file
        if facies is not None:
            las_exporter = LASExporter()
            interpreted_curves = {"FACIES": facies}

            if confidence_scores is not None:
                interpreted_curves["CONFIDENCE"] = confidence_scores

            las_file = output_path / "interpreted_logs" / f"{well_name}.las"
            las_exporter.export_with_interpretation(
                str(las_file),
                depth,
                original_curves,
                interpreted_curves,
                well_header,
                curve_descriptions={
                    "FACIES": "Interpreted Facies Code (ML)",
                    "CONFIDENCE": "Prediction Confidence (0-1)",
                },
            )

        # Export formation tops
        if formation_tops:
            tops_file = output_path / "formation_tops" / f"{well_name}_tops.txt"
            FormationTopExporter.export_to_ascii(
                formation_tops,
                str(tops_file),
                well_name,
                well_header.get("uwi") if well_header else None,
            )

        # Export QC/confidence report
        if confidence_scores is not None:
            qc_data = pd.DataFrame(
                {
                    "Depth_m": depth,
                    "Confidence": confidence_scores,
                    "Needs_Review": confidence_scores < 0.5,
                }
            )
            qc_file = output_path / "qc" / f"{well_name}_confidence.csv"
            qc_data.to_csv(qc_file, index=False)

        # Create README
        readme = output_path / "README.txt"
        with open(readme, "w") as f:
            f.write("PyGeomodeling Interpretation Package\n")
            f.write("=" * 50 + "\n\n")
            f.write(f"Well: {well_name}\n")
            f.write(
                f"Export Date: "
                f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n"
            )
            f.write("Contents:\n")
            f.write(
                "- interpreted_logs/: LAS files with ML-interpreted facies curves\n"
            )
            f.write("- formation_tops/: Formation boundary picks\n")
            f.write("- qc/: Quality control and confidence metrics\n\n")
            f.write("Import Instructions:\n")
            f.write("1. Import LAS files from interpreted_logs/ folder\n")
            f.write("2. Load formation tops from formation_tops/ folder\n")
            f.write("3. Review low-confidence intervals flagged in qc/ reports\n")

        logger.info("Exported interpretation package to %s", output_dir)


def create_correction_template(
    well_name: str,
    depth: np.ndarray,
    predictions: np.ndarray,
    confidence: np.ndarray,
    output_file: str,
):
    """
    Create template file for expert corrections.

    Experts can edit this file and re-import corrections for model retraining.

    Args:
        well_name: Well identifier
        depth: Depth array
        predictions: Model predictions
        confidence: Prediction confidence
        output_file: Path to output CSV template
    """
    df = pd.DataFrame(
        {
            "Well": well_name,
            "Depth_m": depth,
            "ML_Prediction": predictions,
            "Confidence": confidence,
            "Expert_Correction": predictions,  # Pre-fill with predictions
            "Corrected": False,  # Flag for whether expert reviewed
            "Notes": "",  # Free-text notes
        }
    )

    # Filter to show low-confidence predictions first
    df = df.sort_values("Confidence")

    df.to_csv(output_file, index=False)
    logger.info("Created correction template: %s", output_file)
    logger.info("  - Edit 'Expert_Correction' column to fix errors")
    logger.info("  - Set 'Corrected' to TRUE for reviewed samples")
    logger.info("  - Add notes in 'Notes' column as needed")


def import_expert_corrections(
    correction_file: str,
) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Import expert corrections for model retraining.

    Args:
        correction_file: Path to edited correction template

    Returns:
        corrected_data: DataFrame with corrected labels
        original_data: DataFrame with original predictions
    """
    df = pd.read_csv(correction_file)

    # Filter to corrected samples
    corrected = df[df["Corrected"]].copy()

    if len(corrected) == 0:
        raise ValueError(
            "No corrected samples found in file",
            "Set 'Corrected' column to TRUE for reviewed samples",
        )

    logger.info(
        "Imported %d expert corrections from %s", len(corrected), correction_file
    )

    return corrected, df
