"""
铁塔绘图模块 (Tower Drawing Module)
使用 matplotlib 绘制铁塔立面图、截面图和内力图。
"""

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
import os

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'WenQuanYi Zen Hei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False


class TowerDrawer:
    """铁塔绘图器"""

    def __init__(self, calc_results, params):
        self.results = calc_results
        self.params = params
        self.tower_height = params.get("tower_height", 30)
        self.tower_type = params.get("tower_type", "自立塔")
        self.base_width = calc_results.get("塔底宽度", 4.2)
        self.top_width = calc_results.get("塔顶宽度", 1.2)

    def draw_elevation(self, save_path=None):
        """绘制铁塔立面图"""
        fig, ax = plt.subplots(1, 1, figsize=(10, 14))
        fig.suptitle('铁塔立面图', fontsize=18, fontweight='bold', y=0.98)

        H = self.tower_height
        bw = self.base_width
        tw = self.top_width

        x_left = np.linspace(-bw/2, -tw/2, 50)
        x_right = np.linspace(bw/2, tw/2, 50)
        y = np.linspace(0, H, 50)
        ax.plot(x_left, y, 'b-', linewidth=2.5, label='主材')
        ax.plot(x_right, y, 'b-', linewidth=2.5)

        n_horiz = 10
        for i in range(n_horiz + 1):
            yi = H * i / n_horiz
            xi_left = -bw/2 + (bw/2 - tw/2) * (yi / H)
            xi_right = bw/2 - (bw/2 - tw/2) * (yi / H)
            ax.plot([xi_left, xi_right], [yi, yi], 'g-', linewidth=1.0, alpha=0.7)

        for i in range(n_horiz):
            y1 = H * i / n_horiz
            y2 = H * (i + 1) / n_horiz
            x1l = -bw/2 + (bw/2 - tw/2) * (y1 / H)
            x1r = bw/2 - (bw/2 - tw/2) * (y1 / H)
            x2l = -bw/2 + (bw/2 - tw/2) * (y2 / H)
            x2r = bw/2 - (bw/2 - tw/2) * (y2 / H)
            ax.plot([x1l, x2r], [y1, y2], 'r-', linewidth=1.2, alpha=0.6)
            ax.plot([x1r, x2l], [y1, y2], 'r-', linewidth=1.2, alpha=0.6)

        z_wire = self.results.get("导线高度", H * 0.92)
        arm_length = bw * 0.8
        ax.plot([-arm_length, arm_length], [z_wire, z_wire], 'k-', linewidth=3, label='导线横担')
        ax.plot([-arm_length, -tw/2], [z_wire, z_wire - 1.5], 'k-', linewidth=1.5)
        ax.plot([arm_length, tw/2], [z_wire, z_wire - 1.5], 'k-', linewidth=1.5)
        for x_arm in [-arm_length * 0.85, 0, arm_length * 0.85]:
            ax.plot([x_arm, x_arm], [z_wire, z_wire - 2], 'k--', linewidth=1)
            ax.plot(x_arm, z_wire - 2, 'ko', markersize=5)

        z_ground = self.results.get("地线高度", H * 0.98)
        ax.plot([-tw/2, 0], [z_wire, z_ground], 'm-', linewidth=2, label='地线支架')
        ax.plot([tw/2, 0], [z_wire, z_ground], 'm-', linewidth=2)
        ax.plot(0, z_ground, 'm^', markersize=8)

        foundation_b = self.results.get("基础边长", 3.0)
        foundation_d = self.results.get("基础埋深", 2.0)
        rect = patches.Rectangle((-foundation_b/2, -foundation_d), foundation_b, foundation_d,
                                  linewidth=2, edgecolor='brown', facecolor='tan', alpha=0.5)
        ax.add_patch(rect)
        ax.text(0, -foundation_d - 0.5, f'基础 {foundation_b:.2f}m × {foundation_b:.2f}m\n埋深 {foundation_d:.1f}m',
                ha='center', fontsize=9, color='brown')

        ax.annotate(f'塔高 {H:.1f}m', xy=(bw/2 + 0.5, H/2), fontsize=11,
                    arrowprops=dict(arrowstyle='<->', color='blue'), xytext=(bw/2 + 2.5, H/2))
        ax.annotate(f'底宽 {bw:.2f}m', xy=(0, 0.3), fontsize=10, ha='center', color='blue')

        ax.plot([bw/2 + 0.3, bw/2 + 0.3], [0, H], 'k-', linewidth=0.5)
        ax.plot([bw/2 + 0.2, bw/2 + 0.4], [0, 0], 'k-', linewidth=0.5)
        ax.plot([bw/2 + 0.2, bw/2 + 0.4], [H, H], 'k-', linewidth=0.5)

        ax.set_xlim(-arm_length - 2, arm_length + 5)
        ax.set_ylim(-foundation_d - 2, H + 2)
        ax.set_aspect('equal')
        ax.grid(True, alpha=0.3)
        ax.legend(loc='upper left', fontsize=9)
        ax.set_xlabel('宽度 (m)', fontsize=11)
        ax.set_ylabel('高度 (m)', fontsize=11)

        info_text = (
            f"塔型: {self.tower_type}\n"
            f"塔高: {H:.1f} m\n"
            f"抗风级别: {self.params.get('wind_level', 10)} 级\n"
            f"材质: {self.params.get('material', 'Q345')}\n"
            f"电压: {self.params.get('voltage', '110kV')}\n"
            f"主材: {self.results.get('主材规格', 'N/A')}"
        )
        ax.text(0.02, 0.98, info_text, transform=ax.transAxes, fontsize=9,
                verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))

        plt.tight_layout()
        if save_path:
            plt.savefig(save_path, dpi=150, bbox_inches='tight')
            plt.close()
            return save_path
        plt.close()
        return fig

    def draw_cross_section(self, save_path=None):
        """绘制铁塔截面图"""
        fig, axes = plt.subplots(1, 2, figsize=(14, 6))
        fig.suptitle('铁塔截面图', fontsize=16, fontweight='bold')

        bw = self.base_width
        tw = self.top_width

        for ax, w, title in [(axes[0], bw, f'底截面 (宽 {bw:.2f}m)'),
                              (axes[1], tw, f'顶截面 (宽 {tw:.2f}m)')]:
            hw = w / 2
            corners = [(-hw, -hw), (hw, -hw), (hw, hw), (-hw, hw)]
            for i in range(4):
                p1 = corners[i]
                p2 = corners[(i + 1) % 4]
                ax.plot([p1[0], p2[0]], [p1[1], p2[1]], 'b-', linewidth=2.5)
            ax.plot([-hw, hw], [-hw, hw], 'r--', linewidth=1, alpha=0.5)
            ax.plot([-hw, hw], [hw, -hw], 'r--', linewidth=1, alpha=0.5)
            for i, (cx, cy) in enumerate(corners):
                ax.plot(cx, cy, 'bs', markersize=8)
                ax.annotate(f'角钢{i+1}', (cx, cy), textcoords="offset points",
                            xytext=(8, 8), fontsize=8)
            ax.set_xlim(-hw - 0.5, hw + 0.5)
            ax.set_ylim(-hw - 0.5, hw + 0.5)
            ax.set_aspect('equal')
            ax.grid(True, alpha=0.3)
            ax.set_title(title, fontsize=12)

        plt.tight_layout()
        if save_path:
            plt.savefig(save_path, dpi=150, bbox_inches='tight')
            plt.close()
            return save_path
        plt.close()
        return fig

    def draw_internal_force(self, save_path=None):
        """绘制内力图"""
        forces = self.results.get("内力分段", [])
        if not forces:
            return None

        fig, axes = plt.subplots(1, 3, figsize=(18, 6))
        fig.suptitle('塔身内力分布图', fontsize=16, fontweight='bold')

        z = [f["z"] for f in forces]
        M = [f["M"] for f in forces]
        V = [f["V"] for f in forces]
        N = [f["N"] for f in forces]

        axes[0].plot(M, z, 'b-', linewidth=2)
        axes[0].fill_betweenx(z, 0, M, alpha=0.2, color='blue')
        axes[0].set_xlabel('弯矩 M (kN·m)', fontsize=11)
        axes[0].set_ylabel('高度 z (m)', fontsize=11)
        axes[0].set_title('弯矩图', fontsize=13)
        axes[0].grid(True, alpha=0.3)
        axes[0].invert_yaxis()

        axes[1].plot(V, z, 'r-', linewidth=2)
        axes[1].fill_betweenx(z, 0, V, alpha=0.2, color='red')
        axes[1].set_xlabel('剪力 V (kN)', fontsize=11)
        axes[1].set_title('剪力图', fontsize=13)
        axes[1].grid(True, alpha=0.3)
        axes[1].invert_yaxis()

        axes[2].plot(N, z, 'g-', linewidth=2)
        axes[2].fill_betweenx(z, 0, N, alpha=0.2, color='green')
        axes[2].set_xlabel('轴力 N (kN)', fontsize=11)
        axes[2].set_title('轴力图', fontsize=13)
        axes[2].grid(True, alpha=0.3)
        axes[2].invert_yaxis()

        plt.tight_layout()
        if save_path:
            plt.savefig(save_path, dpi=150, bbox_inches='tight')
            plt.close()
            return save_path
        plt.close()
        return fig

    def draw_all(self, output_dir="output"):
        """绘制全部图纸"""
        os.makedirs(output_dir, exist_ok=True)
        paths = {}
        paths["立面图"] = self.draw_elevation(os.path.join(output_dir, "tower_elevation.png"))
        paths["截面图"] = self.draw_cross_section(os.path.join(output_dir, "tower_section.png"))
        paths["内力图"] = self.draw_internal_force(os.path.join(output_dir, "tower_forces.png"))
        return paths