import dev.axllm.ax.*;
import java.nio.file.*;
import java.util.*;

public final class AxgenOptimizationExample {
  static String apiKey() {
    String apiKey = System.getenv("OPENAI_API_KEY");
    if (apiKey == null || apiKey.isBlank()) apiKey = System.getenv("OPENAI_APIKEY");
    if (apiKey == null || apiKey.isBlank()) {
      throw new IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.");
    }
    return apiKey;
  }

  static AxAIService client() {
    return Ax.ai("openai",
        Map.of("api_key", apiKey(), "model", System.getenv().getOrDefault("AX_OPENAI_MODEL", "gpt-6-luna"), "model_config", Map.of("temperature", 0.0)));
  }

  static final class ExampleOptimizer implements OptimizerEngine {
    public String name() { return "example"; }
    public String version() { return "1"; }
    public Map<String, Object> optimize(Map<String, Object> request) {
      return Map.of("componentMap", Map.of("priority::instruction", "Classify operational risk. Use high for production-impacting urgency."), "metadata", Map.of("source", "axgen"));
    }
  }

  public static void main(String[] args) throws Exception {
    AxGen program = new AxGen(Ax.s("emailText:string -> priority:class \"high, normal, low\", rationale:string"), Map.of("id", "priority", "instruction", "Classify the email priority."));
    Map<String, Object> baseline = program.forward(client(), Map.of("emailText", "Production checkout is failing for enterprise customers."));
    Map<String, Object> artifact = program.optimizeWith(new ExampleOptimizer(), List.of(Map.of("emailText", "URGENT: checkout is down", "priority", "high")), Map.of("apply", false));
    program.applyOptimization(Json.stringify(artifact));
    Map<String, Object> after = program.forward(client(), Map.of("emailText", "Production checkout is failing for enterprise customers."));
    double score = Ax.novelF1Score("", String.valueOf(after.get("rationale")), "Production checkout is failing for enterprise customers.");
    System.out.println(Json.stringify(Map.of("baseline", baseline, "after", after, "rationaleNovelF1", score)));
  }
}

import dev.axllm.ax.*;
import java.util.*;

public final class GepaOptimizationExample {
  static String apiKey() {
    String apiKey = System.getenv("OPENAI_API_KEY");
    if (apiKey == null || apiKey.isBlank()) apiKey = System.getenv("OPENAI_APIKEY");
    if (apiKey == null || apiKey.isBlank()) {
      throw new IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.");
    }
    return apiKey;
  }

  static AxAIService client() {
    return Ax.ai("openai",
        Map.of("api_key", apiKey(), "model", System.getenv().getOrDefault("AX_OPENAI_MODEL", "gpt-6-luna")));
  }

  static Map<String, Object> task(String text, String priority) {
    return Map.of("input", Map.of("emailText", text), "expectedOutput", Map.of("priority", priority));
  }

  @SuppressWarnings("unchecked")
  static Map<String, Object> map(Object value) {
    return value instanceof Map<?, ?> ? (Map<String, Object>) value : Map.of();
  }

  static final class PriorityEvaluator implements OptimizerEvaluator {
    final AxGen program;
    final AxAIService client;
    PriorityEvaluator(AxGen program, AxAIService client) { this.program = program; this.client = client; }
    public Map<String, Object> evaluate(Map<String, Object> candidateMap, Map<String, Object> options) {
      var result = program.evaluateOptimization(client, options.get("dataset"), candidateMap, options);
      double sum = 0;
      var rows = (List<?>) result.get("rows");
      for (Object item : rows) {
        var row = map(item);
        var expected = map(map(row.get("input")).get("expectedOutput"));
        var output = map(map(row.get("prediction")).get("output"));
        double score = Objects.equals(output.get("priority"), expected.get("priority")) ? 1 : 0;
        row.put("scores", Map.of("accuracy", score));
        row.put("scalar", score);
        sum += score;
      }
      result.put("avg", sum / rows.size());
      return result;
    }
  }

  public static void main(String[] args) throws Exception {
    AxAIService student = client();
    AxGen program = Ax.ax("emailText:string -> priority:class \"high, normal, low\", rationale:string", Map.of("id", "priority", "instruction", "Classify the email priority."), null);
    Map<String, Object> baseline = program.forward(student, Map.of("emailText", "Production checkout is failing for enterprise customers."));
    Map<String, Object> request = Map.of(
        "programKind", "axgen", "components", program.getOptimizableComponents(),
        "dataset", Map.of("train", List.of(task("URGENT: checkout is down", "high"), task("Weekly product newsletter", "low")),
            "validation", List.of(task("Production payments are failing", "high"), task("Optional community meetup invite", "low"))),
        "options", Map.of("numTrials", 1, "minibatchSize", 2, "maxMetricCalls", 12, "bootstrap", false, "seed", 7));
    var artifact = new AxGEPA(student, Map.of("seed", 7)).optimize(request, new PriorityEvaluator(program, student));
    double original = ((Number) artifact.get("baselineScore")).doubleValue();
    double best = ((Number) map(artifact.get("metadata")).get("bestScore")).doubleValue();
    System.out.println(Json.stringify(Map.of("baselinePrediction", baseline, "baselineScore", original,
        "baselineScores", artifact.get("baselineScores"), "optimizedScore", best, "delta", best - original)));
  }
}

import dev.axllm.ax.*;
import java.nio.file.*;
import java.util.*;

public final class ArtifactOptimizationExample {
  static String apiKey() {
    String apiKey = System.getenv("OPENAI_API_KEY");
    if (apiKey == null || apiKey.isBlank()) apiKey = System.getenv("OPENAI_APIKEY");
    if (apiKey == null || apiKey.isBlank()) {
      throw new IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.");
    }
    return apiKey;
  }

  static AxAIService client() {
    return Ax.ai("openai",
        Map.of("api_key", apiKey(), "model", System.getenv().getOrDefault("AX_OPENAI_MODEL", "gpt-5.4-mini"), "model_config", Map.of("temperature", 0.0)));
  }

  static final class ExampleOptimizer implements OptimizerEngine {
    public String name() { return "example"; }
    public String version() { return "1"; }
    public Map<String, Object> optimize(Map<String, Object> request) {
      return Map.of("componentMap", Map.of("priority::instruction", "Classify operational risk. Use high for production-impacting urgency."), "metadata", Map.of("source", "artifact"));
    }
  }

  public static void main(String[] args) throws Exception {
    AxGen program = new AxGen(Ax.s("emailText:string -> priority:class \"high, normal, low\", rationale:string"), Map.of("id", "priority", "instruction", "Classify the email priority."));
    Map<String, Object> baseline = program.forward(client(), Map.of("emailText", "Production checkout is failing for enterprise customers."));
    Map<String, Object> artifact = program.optimizeWith(new ExampleOptimizer(), List.of(Map.of("emailText", "URGENT: checkout is down", "priority", "high")), Map.of("apply", false));
    program.applyOptimization(Json.stringify(artifact));
    Map<String, Object> after = program.forward(client(), Map.of("emailText", "Production checkout is failing for enterprise customers."));
    System.out.println(Json.stringify(Map.of("baseline", baseline, "after", after)));
  }
}

import dev.axllm.ax.*;
import dev.axllm.ax.runtime.quickjs.*;
import java.util.*;
import java.util.function.Consumer;

public final class AgentPlaybookEvolveExample {
  static String apiKey() {
    String value = System.getenv("OPENAI_API_KEY");
    if (value == null || value.isBlank()) value = System.getenv("OPENAI_APIKEY");
    if (value == null || value.isBlank()) {
      throw new IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.");
    }
    return value;
  }

  public static void main(String[] args) throws Exception {
    AxAIService client = Ax.ai("openai", Map.of(
        "api_key", apiKey(),
        "model", System.getenv().getOrDefault("AX_OPENAI_MODEL", "gpt-6-luna")));

    Map<String, Object> bullet = Map.of(
        "id", "failures-to-avoid-00001",
        "section", "failures_to_avoid",
        "content", "Check the available evidence before answering.",
        "helpfulCount", 0,
        "harmfulCount", 0,
        "createdAt", "2026-07-15T00:00:00.000Z",
        "updatedAt", "2026-07-15T00:00:00.000Z");
    Map<String, Object> seed = Map.of(
        "playbook", Map.of(
            "version", 1,
            "sections", Map.of("failures_to_avoid", List.of(bullet)),
            "updatedAt", "2026-07-15T00:00:00.000Z"),
        "artifact", Map.of("feedback", List.of(), "history", List.of()));

    List<Object> observedCitations = new ArrayList<>();
    List<Object> playbookUpdates = new ArrayList<>();
    Consumer<List<Object>> citationObserver = observedCitations::add;
    Consumer<Map<String, Object>> playbookObserver = playbookUpdates::add;
    AxAgent assistant = Ax.agent(
        "question:string -> answer:string",
        Map.of(
            "ai", client,
            "contextFields", List.of(),
            "runtime", Map.of("language", "JavaScript"),
            "playbook", Map.of("playbook", seed, "seed", 7, "onUpdate", playbookObserver),
            "citations", Map.of("surface", "hidden", "onCitations", citationObserver)));
    assistant
        .setInstruction("Answer from evidence and state uncertainty plainly.")
        .addActorInstruction("Before finishing, verify the answer against the collected evidence.");

    try (AxQuickJsCodeRuntime runtime = new AxQuickJsCodeRuntime()) {
      Map<String, Object> answer = assistant.forward(
          client,
          Map.of("question", "What should a support agent verify before answering?"),
          Map.of("runtime", runtime, "max_actor_steps", 8));

      Map<String, Object> dataset = Map.of(
          "train", List.of(Map.of(
              "input", Map.of("question", "Give a concise evidence-first answer."),
              "score", 0)));
      Map<String, Object> evolution = assistant.playbook(null).evolve(
          dataset,
          Map.of("verify", true, "maxProposals", 1, "runtime", runtime));

      // Each train task replays on the runtime. A replay that fails scores 0
      // and would be mined as the weakness, so stop instead.
      List<Object> replays = new java.util.ArrayList<>();
      for (Object raw : (List<?>) evolution.get("records")) {
        Map<?, ?> record = (Map<?, ?>) raw;
        Map<?, ?> prediction = record.get("prediction") instanceof Map<?, ?> found ? found : Map.of();
        if (record.get("error") != null || "error".equals(prediction.get("completionType"))) {
          throw new IllegalStateException("evolve replay failed: " + (record.get("error") != null ? record.get("error") : prediction.get("error")));
        }
        replays.add(prediction.get("completionType"));
      }

      System.out.println(Json.pretty(answer));
      System.out.println("replays: " + replays);
      System.out.println("citations: " + (observedCitations.isEmpty() ? List.of() : observedCitations.get(observedCitations.size() - 1)));
      System.out.println("run-end updates: " + playbookUpdates.size());
      System.out.println("outcomes: " + evolution.get("outcomes"));
      System.out.println(assistant.getPlaybook().render());
    }
  }
}