from arize.exporter import ArizeExportClient
from arize.utils.types import Environments
from datetime import datetime, timedelta, timezone

client = ArizeExportClient()

primary_df = client.export_model_to_df(
    space_id = "YOUR_ARIZE_SPACE_ID",
    model_id = "YOUR_ARIZE_MODEL_ID",
    environment = Environments.TRACING,
    start_time = datetime.now(timezone.utc) - timedelta(days=7),
    end_time   = datetime.now(timezone.utc),
)

import pandas as pd

# Example: Aggregate all outputs in a trace
trace_df = (
    primary_df.groupby("context.trace_id")
      .agg({
          "attributes.input.value": "first",
          "attributes.output.value": lambda x: " ".join(x.dropna()),
          # Add other aggregations as needed
      })
)

TRACE_EVAL_PROMPT = """
You are evaluating the overall quality and correctness of an LLM application's response to a user request.

You will be given:
1. The user input that initiated the trace
2. The full output(s) generated during the trace

##
User Input:
{attributes.input.value}

Trace Output:
{attributes.output.value}
##

Respond with exactly one word: `correct` or `incorrect`.
- `correct` → the trace achieves the intended goal.
- `incorrect` → the trace fails to achieve the goal or is low quality.
"""

from phoenix.evals import create_classifier
from phoenix.evals.evaluators import async_evaluate_dataframe
from phoenix.evals.llm import LLM

llm = LLM(provider="openai", model="gpt-4o-mini")

correctness_evaluator = create_classifier(
    name="correctness",
    llm=llm,
    prompt_template=TRACE_EVAL_PROMPT,
    choices={"correct": 1.0, "incorrect": 0.0},
)

results_df = await async_evaluate_dataframe(
    dataframe=trace_df,
    evaluators=[correctness_evaluator],
)

from arize.pandas.logger import Client
from phoenix.evals.utils import to_annotation_dataframe
import ast

import pandas as pd
client = Client()

root_spans = primary_df[primary_df["parent_id"].isna()][["context.trace_id", "context.span_id"]]

# Merge results with root spans to align on trace_id
results_with_spans = pd.merge(
    results_df.reset_index(), root_spans, on="context.trace_id", how="left"
).set_index("context.span_id", drop=False)

# Format for logging
correctness_eval_df = to_annotation_dataframe(results_with_spans)

# using trace_eval prefix to rename columns
correctness_eval_df = correctness_eval_df.rename(columns={
    "label": "trace_eval.correctness.label",
    "score": "trace_eval.correctness.score",
    "explanation": "trace_eval.correctness.explanation"
})

client.log_evaluations_sync(correctness_eval_df, 'your-project-name')