from sqlalchemy import Column, Integer, Float
from sqlalchemy import ForeignKeyConstraint
from app.database import Base  # 统一的基础模型


class FluxCdInputData(Base):
    """
    通量镉预测模型输入数据模型

    收集计算土壤镉通量平衡所需的输入参数：
      - 初始土壤镉总量
      - 四大输入通量(大气沉降、灌溉水、农业投入品、其他)
      - 四大输出通量(地下渗漏、地表径流、籽粒移除、秸秆移除)

    注：所有通量单位均为g/ha/a
    """
    __tablename__ = 'FluxCd_input_data'

    # 主键配置
    id = Column('ID', Integer, primary_key=True, autoincrement=True,
                comment='自增主键')

    # 外键字段 - 关联农田样点
    farmland_id = Column('Farmland_ID', Integer, nullable=False,
                         comment='关联农地ID')
    sample_id = Column('Sample_ID', Integer, nullable=False,
                       comment='关联样点ID')

    # 初始土壤镉总量
    initial_cd = Column('Initial_Cd', Float, nullable=False,
                        comment='土壤初始Cd总量[g/ha](100m×100m栅格)')

    # 输入通量
    atmospheric_deposition = Column('DQCJ_Cd', Float, nullable=False,
                                    comment='大气沉降输入镉[g/ha/a]')
    irrigation_input = Column('GGS_Cd', Float, nullable=False,
                              comment='灌溉水输入镉[g/ha/a]')
    agro_chemicals_input = Column('NCP_Cd', Float, nullable=False,
                                  comment='农业投入品输入镉[g/ha/a]')

    # 输出通量 (带默认值)
    groundwater_leaching = Column('DX_Cd', Float, default=0.023,
                                  comment='地下渗漏输出镉[g/ha/a]')
    surface_runoff = Column('DB_Cd', Float, default=0.368,
                            comment='地表径流输出镉[g/ha/a]')
    grain_removal = Column('ZL_Cd', Float, nullable=False,
                           comment='籽粒移除输出镉[g/ha/a]')
    straw_removal = Column('JG_Cd', Float, nullable=False,
                           comment='秸秆移除输出镉[g/ha/a]')

    # 复合外键约束
    __table_args__ = (
        ForeignKeyConstraint(
            ['Farmland_ID', 'Sample_ID'],
            ['Farmland_data.Farmland_ID', 'Farmland_data.Sample_ID']
        ),
        {'comment': '通量镉预测模型输入数据'}
    )

    # 计算输入总通量
    def input_flux(self):
        """计算镉输入总通量[g/ha/a]"""
        return self.atmospheric_deposition + self.irrigation_input + self.agro_chemicals_input

    # 计算输出总通量
    def output_flux(self):
        """计算镉输出总通量[g/ha/a]"""
        return self.groundwater_leaching + self.surface_runoff + self.grain_removal + self.straw_removal

    # 计算净通量
    def net_flux(self):
        """计算镉净通量[g/ha/a]（输入-输出）"""
        return self.input_flux() - self.output_flux()