from gdisdk.connectors.gdimConnector import log_in
from gdisdk.modules.filters import DropDuplicateRows
from gdisdk.modules.operators import AddTableColumns
from gdisdk.modules.readers import CsvReader
from gdisdk.modules.writers import GdimTableWriter
from gdisdk.pipeline.pipeline import PipeLine

# ---- 实例化模块 -------
# 读取csv数据
read_csv = CsvReader(mname="ReadCsv")
read_csv.file = "地层表.csv"
read_csv.dtype = {"bore_number": "str"}  # 设置钻孔编号列为字符串类型

# 去除钻孔编号重复的行
drop_duplicate = DropDuplicateRows(mname="DropDuplicate")
drop_duplicate.subset = ["bore_number"]  # 设置需要去重的列名，支持设置多列

# 增加一列，用于存储新的地层名称 - 地层编号+岩性名称
add_column = AddTableColumns(mname="AddColumn")
add_column.column_names = [{"name": "layer_name", "title": "layer_name", "unit": "-"}]  # 设置新列的名称、标题、单位
add_column.column_templates = {"layer_name": "{layer_number} + {material_name}"}

# ----- 新增模块 ------
# 写入GDIM数据模块
write_table = GdimTableWriter(mname="WriteTable")
write_table.table_names = "钻孔一览表"
write_table.fields_mapping = {
    "bore_number": "钻孔编号",
    "bore_top": "孔顶标高",
    "layer_number": "地层编号",
    "material_name": "岩性名称",
    "layer_name": "地层名称",
}
# ----- 新增模块 ------

# ----- 组织pipeline -------
# 初始化pipeline
pipeline = PipeLine(app_name="GetBoreTable", app_title="获取地层表")

# --- 新增代码 ----
pipeline.update_gdim_state(
    token=log_in(user_name="您的账号", password="您的密码"),
    proj_id="1967572096104902658",
)
# --- 新增代码 ----

# ---- 更新代码 ----
# 链接模块
links = (
    read_csv.OutputTable >> drop_duplicate.InputTable
    | drop_duplicate.OutputTable >> add_column.InputTable
    | add_column.OutputTable >> write_table.InputData
)
# ---- 更新代码 ----

# 添加模块和链接至pipeline
pipeline.add_links(links)

# pipeline.draw_pipeline(view=True)

# 运行pipeline
pipeline.run()
print("成功写入数据至GDIM！")