from mrra.data.trajectory import TrajectoryBatch
from mrra.data.activity import ActivityExtractor
from mrra.analysis.activity_purpose import ActivityPurposeAssigner
from mrra.graph.mobility_graph import MobilityGraph, GraphConfig
from mrra.retriever.graph_rag import GraphRAGGenerate
from mrra.graph.pattern import PatternGenerate
from mrra.agents.builder import build_mrra_agent
from mrra.agents.subagents import make_llm
from mrra.persist.cache import CacheManager, compute_tb_hash

import pandas as pd

# 1) 加载轨迹数据（必须包含列：user_id, timestamp, latitude, longitude）
df = pd.read_csv("your_traj.csv")
tb = TrajectoryBatch(df)  # 自动解析时间，添加派生列如hour/dow

# 2) 活动提取（radius/grid方法）和LLM并发细化"活动目的"
ext_cfg = dict(
    method="radius", 
    radius_m=300, 
    min_dwell_minutes=30, 
    max_gap_minutes=90, 
    grid_size_m=200
)
acts = ActivityExtractor(tb, **ext_cfg).extract()

llm_cfg = dict(
    provider='openai-compatible', 
    model='qwen-plus', 
    base_url='https://dashscope.aliyuncs.com/compatible-mode/v1', 
    api_key='YOUR_API_KEY', 
    temperature=0.2
)
llm = make_llm(**llm_cfg)  # 注意：这是一个可调用客户端，不是字典
acts = ActivityPurposeAssigner(tb, llm=llm, concurrency=8).assign(acts)

# 3) 可选：缓存结果以便重用
cm = CacheManager()
tb_hash = compute_tb_hash(tb)
cm.save_activities(tb_hash, "default", acts)

# 4) 构建移动图（注入"目的分配"活动以避免在图构建期间重新调用LLM）
cfg = GraphConfig(grid_size_m=200, min_dwell_minutes=5, use_activities=True)
mg = MobilityGraph(tb, cfg, activities=acts, assume_purposes_assigned=True)

# 5) 图检索器（GraphRAG），支持目的作为种子
retriever = GraphRAGGenerate(tb=tb, mobility_graph=mg)

# 6) 多智能体反思智能体（两个子智能体 + 置信度投票聚合）
reflection_cfg = dict(
    max_round=1,
    subAgents=[
        {"name": "temporal", "prompt": "从选项中选择最可能的位置id（选择），不要输出坐标。"},
        {"name": "spatial",  "prompt": "从选项中选择最可能的位置id（选择），不要输出坐标。"},
    ],
    aggregator="confidence_weighted_voting",
)
agent = build_mrra_agent(llm=llm_cfg, retriever=retriever, reflection=reflection_cfg)

# 7) 执行三种预测任务
user_id = tb.users()[0]
last_ts = tb.for_user(user_id).iloc[-1]["timestamp_local"].strftime("%Y-%m-%d %H:%M:%S")

for task, extra in (
    ("next_position", {"t": last_ts}),
    ("future_position", {"t": last_ts}),
    ("full_day_traj", {"date": tb.for_user(user_id).iloc[-1]["timestamp_local"].strftime("%Y-%m-%d")}),
):
    payload = {"task": task, "user_id": user_id}
    payload.update(extra)
    res = agent.invoke(payload)
    print(task, res)