from copy import deepcopy
from typing import Any, Dict

from pydantic import Field
from rm_gallery.core.data.schema import DataSample
from rm_gallery.core.reward.base import BaseReward
from rm_gallery.core.reward.composition import SimpleComposition
from rm_gallery.gallery.rm.finance.event_interpretation.event_analysis import EventAnalysisReward
from rm_gallery.gallery.rm.finance.event_interpretation.event_identification import EventIdentificationReward
from rm_gallery.gallery.rm.finance.macro_analysis.concept_explanation import ConceptExplanationReward
from rm_gallery.gallery.rm.finance.macro_analysis.macro_analysis import MacroAnalysisReward
from rm_gallery.gallery.rm.finance.stock_analysis.fundamental_analysis import FundamentalAnalysisReward
from rm_gallery.gallery.rm.finance.stock_analysis.overall_logic import OverallLogicReward
from rm_gallery.gallery.rm.finance.stock_analysis.risk_analysis import RiskAnalysisReward
from rm_gallery.gallery.rm.finance.stock_analysis.valuation_analysis import ValuationAnalysisReward

REWARDS = {
    "concept_explanation": ConceptExplanationReward,
    "macro_analysis": MacroAnalysisReward,
    "overall_logic": OverallLogicReward,

}

WEIGHTS = {
    "concept_explanation": 0.2,
    "macro_analysis": 0.6,
    "overall_logic": 0.2,
}


class MacroAnalysisComposition(SimpleComposition):
    rewards: Dict[str, Dict[str, Any] | BaseReward] = Field(
        default=REWARDS, description="reward modules"
    )
    weights: Dict[str, float] = Field(default=WEIGHTS, description="weight for each reward")
    
    def merge_samples(self, reward_samples: Dict[str, DataSample]):
        final_sample = None
        dimension_scores = {}
        dimension_details = {}
        
        for name,sample in reward_samples.items():
            dimension_scores[name] = {}
            for output in sample.output:
                dimension_scores[name][output.answer.label["model_name"]] = output.answer.reward.score
            if not final_sample:
                final_sample = sample.model_copy(deep=True)
                # Preserve existing metadata instead of overwriting with empty dict
                if final_sample.metadata is None:
                    final_sample.metadata = {}
                dimension_details[name] = deepcopy(sample.metadata) 
                #更新分数 
                for index, output in enumerate(final_sample.output):
                    output.answer.reward.score = output.answer.reward.score * self.weights.get(name,0)         
            else:
                dimension_details[name] = deepcopy(sample.metadata)
                for index, output in enumerate(final_sample.output):
                    output.answer.reward.score += sample.output[index].answer.reward.score * self.weights.get(name,0)
                    output.answer.reward.details.extend(sample.output[index].answer.reward.details)
        dimension_scores["overall_score"] = {}
        for output in final_sample.output:
            dimension_scores["overall_score"][output.answer.label["model_name"]] = output.answer.reward.score
        final_sample.metadata["dimension_scores"] = dimension_scores
        final_sample.metadata["dimension_details"] = dimension_details
        scores = [output.answer.reward.score for output in final_sample.output]
        best_index = scores.index(max(scores))
        final_sample.input[-1].additional_kwargs["best_index"] = best_index
        final_sample.input[-1].additional_kwargs["best_model"] = final_sample.output[best_index].answer.label["model_name"]
        return final_sample