OMOP

def load_omop_table(
    file_path: Path,
    table_name: str,
    *,
    validate: bool | None = None,
    encoding: str | None = None,
) -> pl.DataFrame:
    """Load OMOP CSV into validated Polars DataFrame.

    Args:
        file_path: Path to CSV file.
        table_name: OMOP table name (e.g., 'person', 'condition_occurrence').
        validate: Run schema validation. Uses config default if None.
        encoding: CSV encoding. Uses config default if None.

    Returns:
        Polars DataFrame with validated schema.

    Raises:
        FileNotFoundError: If file doesn't exist.
        ValueError: If validation fails and skip_invalid_rows is False.
    """
    settings = get_omop_settings()
    validate = validate if validate is not None else settings.validate_on_load
    encoding = encoding or settings.csv_encoding

    if not file_path.exists():
        msg = f"OMOP CSV file not found: {file_path}"
        logger.error(msg)
        raise FileNotFoundError(msg)

    start_time = time.time()
    logger.info(
        "Loading OMOP table: %s",
        table_name,
        extra={"file_path": str(file_path), "table_name": table_name},
    )

    try:
        df = pl.read_csv(
            file_path,
            encoding=encoding,
            null_values=["", "NULL"],
            try_parse_dates=True,
            infer_schema_length=settings.infer_schema_length,
        )

        if table_name in CDM_53_RENAMES:
            renames = {k: v for k, v in CDM_53_RENAMES[table_name].items() if k in df.columns}
            if renames:
                df = df.rename(renames)

        duration = time.time() - start_time
        logger.info(
            "Loaded %s rows in %.2fs",
            f"{len(df):,}",
            duration,
            extra={
                "table_name": table_name,
                "row_count": len(df),
                "duration_seconds": duration,
            },
        )

        if validate:
            errors = validate_omop_table_schema(df, table_name)
            if errors:
                error_summary = errors[:5]  # Log first 5 errors
                logger.warning(
                    "Validation found %d issues in %s",
                    len(errors),
                    table_name,
                    extra={
                        "table_name": table_name,
                        "error_count": len(errors),
                        "sample_errors": error_summary,
                    },
                )
                if not settings.skip_invalid_rows:
                    _raise_validation_error(table_name, len(errors))
            # Either no errors or skip_invalid_rows is True - continue to return
        # Return the DataFrame (validated or not)
        return df  # noqa: TRY300

    except Exception:
        logger.exception(
            "Failed to load OMOP table: %s",
            table_name,
            extra={"table_name": table_name, "file_path": str(file_path)},
        )
        raise

def load_omop_dataset(
    data_dir: Path | None = None,
    table_names: list[str] | None = None,
    *,
    validate: bool = True,
) -> dict[str, pl.DataFrame]:
    """Load multiple OMOP tables from directory.

    Args:
        data_dir: Directory containing OMOP CSV files. Uses config default if None.
        table_names: List of table names to load. Loads all CSV files if None.
        validate: Whether to validate schemas during loading.

    Returns:
        Dictionary mapping table names to DataFrames.
    """
    settings = get_omop_settings()
    data_dir = data_dir or settings.raw_data_dir

    if not data_dir.exists():
        msg = f"OMOP data directory not found: {data_dir}"
        logger.error(msg)
        raise FileNotFoundError(msg)

    csv_files = list(data_dir.glob("*.csv"))
    if not csv_files:
        msg = f"No CSV files found in {data_dir}"
        logger.warning(msg)
        return {}

    logger.info(
        "Loading OMOP dataset from %s",
        data_dir,
        extra={"data_dir": str(data_dir), "csv_count": len(csv_files)},
    )

    tables = {}
    for csv_file in csv_files:
        table_name = csv_file.stem.lower()

        if table_names and table_name not in table_names:
            continue

        try:
            tables[table_name] = load_omop_table(csv_file, table_name, validate=validate)
        except (FileNotFoundError, ValueError, OSError):
            logger.exception(
                "Skipping table %s due to error",
                table_name,
                extra={"table_name": table_name},
            )

    logger.info(
        "Loaded %d OMOP tables",
        len(tables),
        extra={"loaded_tables": list(tables.keys())},
    )

    return tables

def load_vocabulary_table(
    file_path: Path,
    table_name: str,
    *,
    validate: bool | None = None,
    encoding: str | None = None,
) -> pl.DataFrame:
    """Load OMOP vocabulary file into validated Polars DataFrame.

    Args:
        file_path: Path to vocabulary file (CSV or TSV).
        table_name: Vocabulary table name (e.g., 'concept', 'vocabulary').
        validate: Run schema validation. Uses config default if None.
        encoding: File encoding. Uses config default if None.

    Returns:
        Polars DataFrame with validated schema.

    Raises:
        FileNotFoundError: If file doesn't exist.
        ValueError: If validation fails and skip_invalid_rows is False.
    """
    settings = get_omop_settings()
    validate = validate if validate is not None else settings.validate_on_load
    encoding = encoding or settings.csv_encoding

    if not file_path.exists():
        msg = f"Vocabulary file not found: {file_path}"
        logger.error(msg, extra={"file_path": str(file_path), "table_name": table_name})
        raise FileNotFoundError(msg)

    start_time = time.time()
    logger.info(
        "Loading vocabulary table: %s",
        table_name,
        extra={"file_path": str(file_path), "table_name": table_name},
    )

    try:
        separator = _detect_separator(file_path, encoding)
        df = pl.read_csv(
            file_path,
            encoding=encoding,
            separator=separator,
            quote_char=None,
            null_values=["", "NULL"],
            try_parse_dates=True,
            infer_schema_length=settings.infer_schema_length,
        )

        duration = time.time() - start_time
        file_size_mb = file_path.stat().st_size / (1024 * 1024)
        logger.info(
            "Loaded %s rows (%.2f MB) in %.2fs",
            f"{len(df):,}",
            file_size_mb,
            duration,
            extra={
                "table_name": table_name,
                "row_count": len(df),
                "file_size_mb": file_size_mb,
                "duration_seconds": duration,
            },
        )

        if validate:
            errors = validate_vocabulary_table_schema(df, table_name)
            if errors:
                error_summary = errors[:5]
                logger.warning(
                    "Validation found %d issues in vocabulary %s",
                    len(errors),
                    table_name,
                    extra={
                        "table_name": table_name,
                        "error_count": len(errors),
                        "sample_errors": error_summary,
                    },
                )
                if not settings.skip_invalid_rows:
                    _raise_validation_error(table_name, len(errors))

        return df  # noqa: TRY300

    except Exception:
        logger.exception(
            "Failed to load vocabulary table: %s",
            table_name,
            extra={"table_name": table_name, "file_path": str(file_path)},
        )
        raise

def load_vocabulary_dataset(
    data_dir: Path | None = None,
    table_names: list[str] | None = None,
    *,
    validate: bool = True,
) -> dict[str, pl.DataFrame]:
    """Load OMOP vocabulary tables from directory.

    Args:
        data_dir: Directory containing vocabulary files. Uses config default if None.
        table_names: List of vocabulary table names to load. Loads all if None.
        validate: Whether to validate schemas during loading.

    Returns:
        Dictionary mapping table names to DataFrames.
    """
    settings = get_omop_settings()
    data_dir = data_dir or settings.vocabulary_dir

    if not data_dir.exists():
        msg = f"Vocabulary directory not found: {data_dir}"
        logger.warning(msg, extra={"data_dir": str(data_dir)})
        return {}

    vocabulary_files = VOCABULARY_FILE_MAPPING
    if table_names:
        include = set(table_names)
        if "concept" in include:
            include.add("concept_cpt4")
        vocabulary_files = {k: v for k, v in vocabulary_files.items() if k in include}

    logger.info(
        "Loading OMOP vocabularies from %s",
        data_dir,
        extra={"data_dir": str(data_dir), "table_count": len(vocabulary_files)},
    )

    tables = {}
    for table_name, file_variants in vocabulary_files.items():
        file_path = _find_vocabulary_file(data_dir, table_name, file_variants)
        if not file_path:
            continue

        try:
            tables[table_name] = load_vocabulary_table(file_path, table_name, validate=validate)
        except (FileNotFoundError, ValueError, OSError):
            logger.exception(
                "Skipping vocabulary table %s due to error",
                table_name,
                extra={"table_name": table_name},
            )

    if "concept_cpt4" in tables:
        if "concept" in tables:
            tables["concept"] = pl.concat([tables["concept"], tables.pop("concept_cpt4")])
            logger.info("Merged concept_cpt4 into concept.")
        else:
            tables.pop("concept_cpt4")

    logger.info(
        "Loaded %d vocabulary tables",
        len(tables),
        extra={"loaded_tables": list(tables.keys())},
    )

    return tables

def get_persons(
    con: duckdb.DuckDBPyConnection | None = None,
    *,
    gender_concept_id: int | None = None,
    year_of_birth_min: int | None = None,
    year_of_birth_max: int | None = None,
    race_concept_id: int | None = None,
    ethnicity_concept_id: int | None = None,
    limit: int | None = None,
) -> pl.DataFrame:
    """Query person records with optional filters."""
    conditions = []
    params = []

    add_filter(conditions, params, "gender_concept_id", "=", gender_concept_id)
    add_filter(conditions, params, "year_of_birth", ">=", year_of_birth_min)
    add_filter(conditions, params, "year_of_birth", "<=", year_of_birth_max)
    add_filter(conditions, params, "race_concept_id", "=", race_concept_id)
    add_filter(conditions, params, "ethnicity_concept_id", "=", ethnicity_concept_id)

    query, params = build_filter_query("ducklake.omop.person", conditions, params, limit=limit)

    with ensure_connection(con) as conn:
        return execute_query(
            conn,
            query,
            params,
            lambda: logger.info("Querying persons", extra={"filter_count": len(conditions)}),
        )

def get_conditions_for_person(
    person_id: int,
    con: duckdb.DuckDBPyConnection | None = None,
    *,
    condition_concept_id: int | None = None,
    start_date: date | None = None,
    end_date: date | None = None,
) -> pl.DataFrame:
    """Query condition occurrences for a person with optional filters."""
    conditions = []
    params = []

    add_filter(conditions, params, "person_id", "=", person_id)
    add_filter(conditions, params, "condition_concept_id", "=", condition_concept_id)
    add_filter(conditions, params, "condition_start_date", ">=", start_date)
    add_filter(conditions, params, "condition_start_date", "<=", end_date)

    query, params = build_filter_query("ducklake.omop.condition_occurrence", conditions, params)

    with ensure_connection(con) as conn:
        return execute_query(
            conn,
            query,
            params,
            lambda: logger.info(
                "Querying conditions for person",
                extra={"person_id": person_id, "filter_count": len(conditions)},
            ),
        )

def get_observations_for_person(
    person_id: int,
    con: duckdb.DuckDBPyConnection | None = None,
    *,
    observation_concept_id: int | None = None,
    start_date: date | None = None,
    end_date: date | None = None,
) -> pl.DataFrame:
    """Query observations for a person with optional filters."""
    conditions = []
    params = []

    add_filter(conditions, params, "person_id", "=", person_id)
    add_filter(conditions, params, "observation_concept_id", "=", observation_concept_id)
    add_filter(conditions, params, "observation_date", ">=", start_date)
    add_filter(conditions, params, "observation_date", "<=", end_date)

    query, params = build_filter_query("ducklake.omop.observation", conditions, params)

    with ensure_connection(con) as conn:
        return execute_query(
            conn,
            query,
            params,
            lambda: logger.info(
                "Querying observations for person",
                extra={"person_id": person_id, "filter_count": len(conditions)},
            ),
        )

def get_visits_for_person(
    person_id: int,
    con: duckdb.DuckDBPyConnection | None = None,
    *,
    visit_concept_id: int | None = None,
    start_date: date | None = None,
    end_date: date | None = None,
) -> pl.DataFrame:
    """Query visit occurrences for a person with optional filters."""
    conditions = []
    params = []

    add_filter(conditions, params, "person_id", "=", person_id)
    add_filter(conditions, params, "visit_concept_id", "=", visit_concept_id)
    add_filter(conditions, params, "visit_start_date", ">=", start_date)
    add_filter(conditions, params, "visit_start_date", "<=", end_date)

    query, params = build_filter_query("ducklake.omop.visit_occurrence", conditions, params)

    with ensure_connection(con) as conn:
        return execute_query(
            conn,
            query,
            params,
            lambda: logger.info(
                "Querying visits for person",
                extra={"person_id": person_id, "filter_count": len(conditions)},
            ),
        )

def get_cohort(
    con: duckdb.DuckDBPyConnection | None = None,
    *,
    has_condition_concept_id: int | None = None,
    has_drug_concept_id: int | None = None,
    min_age: int | None = None,
    max_age: int | None = None,
    gender_concept_id: int | None = None,
) -> pl.DataFrame:
    """Select a cohort of persons matching the given clinical criteria."""
    from_clause = "ducklake.omop.person p"
    where_clauses = []
    params = []

    if has_condition_concept_id is not None:
        from_clause += """
            INNER JOIN ducklake.omop.condition_occurrence co
                ON p.person_id = co.person_id
        """
        add_filter(where_clauses, params, "co.condition_concept_id", "=", has_condition_concept_id)

    if has_drug_concept_id is not None:
        from_clause += """
            INNER JOIN ducklake.omop.drug_exposure de
                ON p.person_id = de.person_id
        """
        add_filter(where_clauses, params, "de.drug_concept_id", "=", has_drug_concept_id)

    current_year = datetime.now(UTC).date().year
    if min_age is not None:
        max_birth_year = current_year - min_age
        add_filter(where_clauses, params, "p.year_of_birth", "<=", max_birth_year)

    if max_age is not None:
        min_birth_year = current_year - max_age
        add_filter(where_clauses, params, "p.year_of_birth", ">=", min_birth_year)

    add_filter(where_clauses, params, "p.gender_concept_id", "=", gender_concept_id)

    where_clause = " AND ".join(where_clauses) if where_clauses else "1=1"
    query = "SELECT DISTINCT p.* FROM " + from_clause + " WHERE " + where_clause

    with ensure_connection(con) as conn:
        return execute_query(
            conn,
            query,
            params,
            lambda: logger.info(
                "Selecting cohort",
                extra={
                    "criteria_count": len(where_clauses),
                    "has_condition": has_condition_concept_id is not None,
                    "has_drug": has_drug_concept_id is not None,
                },
            ),
        )

def get_measurement_values(
    person_id: int,
    measurement_concept_id: int,
    con: duckdb.DuckDBPyConnection | None = None,
    *,
    start_date: date | None = None,
    end_date: date | None = None,
) -> pl.DataFrame:
    """Query measurement values for a person and concept with optional date filters."""
    conditions = []
    params = []

    add_filter(conditions, params, "person_id", "=", person_id)
    add_filter(conditions, params, "measurement_concept_id", "=", measurement_concept_id)
    add_filter(conditions, params, "measurement_date", ">=", start_date)
    add_filter(conditions, params, "measurement_date", "<=", end_date)

    query, params = build_filter_query("ducklake.omop.measurement", conditions, params, order_by="measurement_date")

    with ensure_connection(con) as conn:
        return execute_query(
            conn,
            query,
            params,
            lambda: logger.info(
                "Querying measurement values",
                extra={
                    "person_id": person_id,
                    "measurement_concept_id": measurement_concept_id,
                },
            ),
        )