import math
import paddle
from net import DSSMLayer

class StaticModel():
    def __init__(self, config):
        self.cost = None
        self.config = config
        self._init_hyper_parameters()

    def _init_hyper_parameters(self):
        # 修改超参数初始化
        self.feature_nums = self.config.get("hyper_parameters.feature_nums", [5,5,5,5,5])
        self.embedding_dim = self.config.get("hyper_parameters.embedding_dim", 8)
        self.output_dim = self.config.get("hyper_parameters.output_dim", 16)
        self.hidden_layers = self.config.get("hyper_parameters.hidden_layers", [40, 32])
        self.hidden_acts = self.config.get("hyper_parameters.hidden_acts", ["relu", "relu"])
        self.learning_rate = self.config.get("hyper_parameters.optimizer.learning_rate", 0.001)
        self.margin = self.config.get("hyper_parameters.margin", 0.3)  # 用于损失函数的margin参数
        self.feature_num = len(self.feature_nums)

    def create_feeds(self, is_infer=False):
        # 定义输入数据占位符
        # sample_id = paddle.static.data(
        #    name="sample_id", shape=[-1, 1], dtype='int64')
       
        if not is_infer:
            label = paddle.static.data(
                name="label", shape=[-1, 1], dtype='float32')
            feeds_list.append(label)
                    
        left_features = paddle.static.data(
            name="left_features", shape=[-1, self.feature_num], dtype='float32')
        
        right_features = paddle.static.data(
            name="right_features", shape=[-1, self.feature_num], dtype='float32')

        feeds_list = [left_features, right_features]
        


        return feeds_list

    def net(self, input, is_infer=False):
        # 创建模型实例
        dssm_model = DSSMLayer(
            feature_nums=self.feature_nums,
            embedding_dim=self.embedding_dim,
            output_dim=self.output_dim,
            hidden_layers=self.hidden_layers,
            hidden_acts=self.hidden_acts
        )

    

        if is_infer:
            left_features, right_features = input
        else:
            label,left_features, right_features = input
            paddle.static.Print(left_features, message="lqc left data feature shape:")
            paddle.static.Print(right_features, message="lqc right data feature shape:")
            paddle.static.Print(label, message="lqc label feature shape:")

        # 获取相似度和特征向量
        sim_score, left_vec, right_vec = dssm_model(left_features, right_features)

        self.inference_target_var = sim_score
        self.left_vector = left_vec
        self.right_vector = right_vec

        if is_infer:
            fetch_dict = {
                'similarity': sim_score,
                'left_vector': left_vec,
                'right_vector': right_vec
            }
            return fetch_dict

        # 计算损失
        # 使用带margin的二元交叉熵损失
        pos_mask = paddle.cast(label > 0.5, 'float32')
        neg_mask = 1.0 - pos_mask
        
        positive_loss = -pos_mask * paddle.log(paddle.clip(sim_score, 1e-8, 1.0))
        negative_loss = -neg_mask * paddle.log(paddle.clip(1 - sim_score + self.margin, 1e-8, 1.0))
        
        loss = positive_loss + negative_loss
        avg_cost = paddle.mean(loss)
        
        self._cost = avg_cost

        # 计算accuracy
        predictions = paddle.cast(sim_score > 0.5, 'float32')
        accuracy = paddle.mean(paddle.cast(paddle.equal(predictions, label), 'float32'))

        fetch_dict = {
            'loss': avg_cost,
            'accuracy': accuracy,
            #'similarity': sim_score,
            #'left_vector': left_vec,
            #'right_vector': right_vec
        }
        return fetch_dict

    def create_optimizer(self, strategy=None):
        optimizer = paddle.optimizer.Adam(
            learning_rate=self.learning_rate)
        if strategy is not None:
            import paddle.distributed.fleet as fleet
            optimizer = fleet.distributed_optimizer(optimizer, strategy)
        optimizer.minimize(self._cost)

    def infer_net(self, input):
        return self.net(input, is_infer=True)