from __future__ import annotations

import argparse, os
from pathlib import Path
from datetime import datetime
from numpy import True_
import pandas as pd

from dotenv import dotenv_values, load_dotenv

from feature.doc_alignment import DocAlignmentModel
from evaluator2 import PairwiseRelationshipEvaluator
from io import load_json_rows, save_dataframe
from reporting import write_markdown_report
from neo4j_loader import fetch_rows_from_neo4j
from schema import load_node_schema, load_nodes_schema

load_dotenv()
print("SKIP_FIELDS =", os.getenv("SKIP_FIELDS"))

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Analyze pairwise sample-node property relationships"
    )
    parser.add_argument(
        "--data",
        required=False,
        help="Optional path to JSON array of flattened rows. If omitted, rows are loaded from Neo4j.",
    )
    parser.add_argument(
        "--nodes",
        required=True,
        help="Provide list of nodes to generate data for",
    )
    parser.add_argument(
        "--schema",
        required=False,
        help="Path to YAML schema for the node",
    )
    parser.add_argument(
        "--study",
        dest="study_id",
        required=False,
        help="Study ID for Neo4j filtering",
    )
    parser.add_argument(
        "--output-dir",
        default="outputs",
        help="Directory for analysis outputs",
    )
    parser.add_argument(
        "--neo4j-limit",
        type=int,
        default=None,
        help="Optional row limit when loading from Neo4j",
    )
    return parser

def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    # schema = load_nodes_schema(args.schema, args.nodes)
    nodes = [n.strip() for n in args.nodes.split(",") if n.strip()]
    schemas = load_nodes_schema(nodes)

    for schema in schemas:
        if args.data:
            df = load_json_rows(args.data)
        else:
            df = fetch_rows_from_neo4j(
                schema,
                study_id=args.study_id,
                limit=args.neo4j_limit,
            )
        
        # doc_model = DocAlignmentModel().fit(schema)

        engine = PairwiseRelationshipEvaluator(schema)
        results = engine.evaluate_all_pairs(df)
        
        output_dir = Path(os.environ["OUTPUT_DIR"]) 
        if args.output_dir:
            output_dir = Path(args.output_dir)
            

        output_dir.mkdir(parents=True, exist_ok=True)

        node_name = schema.name
        study_id = args.study_id if args.study_id else "all"
        timestamp = datetime.now().strftime("%Y%m%d_%H%M")
        file_prefix = f"{node_name}_{study_id}_{timestamp}"

        save_dataframe(results, output_dir / node_name / f"{file_prefix}_pairwise_relationships.csv")
        save_dataframe(results, output_dir / node_name /f"{file_prefix}_pairwise_relationships.json")
        write_markdown_report(results, output_dir / node_name / f"{file_prefix}_strong_relationships.md")

    if results.empty:
        print("No pairwise relationships were computed.")
        return

    display_cols = [
        "A",
        "B",
        "classification",
        "strength",
        "predictive_strength",
        "support",
        "determinism",
        "stability",
        "doc_alignment",
        "suffix_match",
        "substring_match",
        "heldout_accuracy",
        "baseline_accuracy",
    ]
    
    print(results[display_cols].head(15).to_string(index=False))
    results = results.rename(columns={"heldout_accuracy": "heldout_accuracy (test)"})
    print(f"\nSaved outputs to {output_dir.resolve()}")

if __name__ == "__main__":
    main()