import numpy as np


# D文件
class DFile:
    def __init__(
        self,
        start_num,
        end_num,
        mileage: np.ndarray,
        center_z,
        left_z,
        right_z,
        tree_height,
    ):
        self.start_num = start_num
        self.end_num = end_num
        self._mileage = mileage
        self._center_z = center_z
        self._left_z = left_z
        self._right_z = right_z
        self._tree_height = tree_height
        pass

    def save(self, out_dir):
        start_num = self.start_num
        end_num = self.end_num
        mileage = self._mileage
        center_z = self._center_z
        left_z = self._left_z
        right_z = self._right_z
        mat = np.zeros((mileage.size, 5))
        mat[0, 0] = start_num
        mat[-1, 0] = end_num
        mat[:, 1] = mileage
        mat[:, 2] = center_z
        mat[:, 3] = left_z
        mat[:, 4] = right_z
        if self._tree_height > 0:
            mat = np.concatenate(
                (
                    np.array([mat[0, :]]),
                    np.array([[197, 0, 0, 0, self._tree_height]]),
                    mat[1:, :],
                )
            )  # 197
            mat = np.concatenate(
                (
                    mat[0:-2, :],
                    np.array([[198, mat[-1,1], 0, 0, self._tree_height]]),
                    np.array([mat[-1, :]]),
                ),
            )  # 198
        with open(f"{out_dir}/D{100+start_num-4000}.txt", "w") as d_file:
            for foo in mat:
                d_file.write(f"{foo[0]:.0f}\t")
                d_file.write(f"{foo[1]:.0f}\t")
                d_file.write(f"{foo[2]}\t")
                d_file.write(f"{foo[3]}\t")
                d_file.write(f"{foo[4]}\t")
                d_file.write("\n")
        with open(f"{out_dir}/X{100+start_num-4000}.txt", "w") as x_file:
            x_file.write("dm\n")
            x_file.write(f"{mileage.size}\n")
            x_file.write("0\n")
            x_file.write(f"{mileage.size*3}\n")
            x_file.write("1\n")
            x_file.write("0\n")
            x_file.write("0\n")
        pass


class ControlFile:
    def __init__(self, segments, tower_start_number, c_file_template_path, out_dir):
        self._c_file_template_path = c_file_template_path
        self._out_dir = out_dir
        self._segments = segments
        self._tower_start_number = tower_start_number
        pass

    def save(self):
        c_file_template_path = self._c_file_template_path
        out_dir = self._out_dir
        segments = self._segments
        tower_start_number = self._tower_start_number
        with open(c_file_template_path) as c_file_template:
            c_template = c_file_template.read()
            c_template = c_template.replace("{SEGMENT_N}", str(len(segments)))
            c_template = c_template.replace(
                "{TOWER_NUM_START}", f"{tower_start_number}"
            )
            c_template = c_template.replace(
                "{SEGMENT}", "\n".join([str(seg - 4000 + 100) for seg in segments])
            )
            c_template = c_template.replace("{SEGMENT_S}", f"{segments[0]-4000+100}")
            c_template = c_template.replace("{SEGMENT_E}", f"{segments[-1]-4000+100}")
            c_template = c_template.replace(
                "{CONTROL_FILE}", f"S{segments[0]-4000+100}"
            )
        # with open(f"{out_dir}/PW{segments[0]-4000+1}.dat", "w") as c_file:
        with open(f"{out_dir}/PW01.dat", "w") as c_file:
            c_file.write(c_template)