from agents import DeepResearchAgent,get_llms
import asyncio
import json
import argparse


if __name__ == '__main__':

    argparser = argparse.ArgumentParser()
    argparser.add_argument("--topic",type=str,help="research topic",default="Using diffusion to generate urban road layout map")
    argparser.add_argument("--anchor_paper_path",type=str,help="PDF path of the anchor paper",default= None)
    argparser.add_argument("--save_file",type=str,default="saves/",help="save file path")
    argparser.add_argument("--improve_cnt",type=int,default= 1,help="experiment refine count")
    argparser.add_argument("--max_chain_length t",type=int,default=5,help="max chain length")
    argparser.add_argument("--min_chain_length",type=int,default=3,help="min chain length")
    argparser.add_argument("--max_chain_numbers",type=int,default=1,help="max chain numbers")
    
    args = argparser.parse_args()

    main_llm , cheap_llm = get_llms()

    topic = args.topic
    anchor_paper_path = args.anchor_paper_path

    deep_research_agent = DeepResearchAgent(llm=main_llm,cheap_llm=cheap_llm,**vars(args))

    print(f"begin to generate idea and experiment of topic {topic}")
    idea,related_experiments,entities,idea_chain,ideas,trend,future,human,year=  deep_research_agent.generate_idea_with_chain(topic)
        
    print(f"succeed to generate idea and experiment of topic {topic}")
    res = {"idea":idea,"related_experiments":related_experiments,"entities":entities,"idea_chain":idea_chain,"ideas":ideas,"trend":trend,"future":future,"year":year,"human":human}
    with open("result.json","w") as f:
        json.dump(res,f)