import shutil
from rscder.gui.actions import ActionManager
from rscder.utils.icons import IconInstance
from rscder.utils.project import Project, PairLayer, ResultPointLayer
from rscder.plugins.basic import BasicPlugin
from PyQt5.QtWidgets import QDialog, QHBoxLayout, QFileDialog, QComboBox, QVBoxLayout, QPushButton, QLabel, QLineEdit, QAction
from PyQt5.QtGui import QIcon
from PyQt5.QtCore import Qt
from osgeo import gdal, gdal_array
import numpy as np
class ExportDialog(QDialog):

    TXT='*.txt'
    TIF='*.tif'

    def __init__(self, parent=None, out_type = TXT):
        super().__init__(parent)

        self.setWindowTitle('Export')
        self.setWindowIcon(IconInstance().LOGO)

        self.out_path = None
        self.result_layer = None

        result_layer_select_label = QLabel('选择结果:')

        result_layer_select = QComboBox(self)

        result_layer_select.addItem('---', None)
        for layer in Project().layers.values():
            for result_layer in layer.layers:
                if isinstance(result_layer, ResultPointLayer):
                    result_layer_select.addItem(IconInstance().VECTOR, result_layer.name, result_layer)
        
        def on_result_layer_select(index):
            self.result_layer = result_layer_select.currentData()

        result_layer_select.currentIndexChanged.connect(on_result_layer_select)

        out_path_label = QLabel('输出路径:')
        out_path_text = QLineEdit(self)
        out_path_text.setReadOnly(True)
        out_path_text.setPlaceholderText('选择输出路径')
        
        def on_out_path_btn():
            select_file = QFileDialog.getSaveFileName(self, '选择输出路径', '',  out_type)
            if select_file[0]:
                out_path_text.setText(select_file[0])
                self.out_path = select_file[0]
        out_path_btn = QPushButton('...', self)
        out_path_btn.clicked.connect(on_out_path_btn)
        
        ok_btn = QPushButton('确定', self)
        ok_btn.clicked.connect(self.accept)
        cancel_btn = QPushButton('取消', self)
        cancel_btn.clicked.connect(self.reject)


        hbox1 = QHBoxLayout()
        hbox1.addWidget(result_layer_select_label)
        hbox1.addWidget(result_layer_select)

        hbox2 = QHBoxLayout()
        hbox2.addWidget(out_path_label)
        hbox2.addWidget(out_path_text)
        hbox2.addWidget(out_path_btn)

        hbox3 = QHBoxLayout()
        hbox3.addWidget(ok_btn,0,alignment=Qt.AlignHCenter)
        hbox3.addWidget(cancel_btn,0,alignment=Qt.AlignHCenter)
        
        vbox = QVBoxLayout()
        vbox.addLayout(hbox1)
        vbox.addLayout(hbox2)
        vbox.addLayout(hbox3)

        self.setLayout(vbox)

class ExportPlugin(BasicPlugin):

    @staticmethod
    def info():
        return {
            'name': 'Export',
            'description': 'Export to other format',
            'author': 'RSCDER',
        }
    
    def set_action(self):
        self.export_txt = QAction(IconInstance().DOCUMENT, '导出为 Arcgis 兼容的TXT', self.mainwindow)
        self.export_txt.triggered.connect(self.export_txt_action)

        self.export_bin = QAction(IconInstance().DOCUMENT, '导出栅格二值变化检测结果')
        self.export_bin.triggered.connect(self.export_bin_action)

        self.export_rgb = QAction(IconInstance().DOCUMENT, '导出变化幅度渲染结果')
        self.export_rgb.triggered.connect(self.export_rgb_action)
        ActionManager().export_menu.addAction(self.export_txt)
        ActionManager().export_menu.addAction(self.export_bin)
        ActionManager().export_menu.addAction(self.export_rgb)

        # self.ctx['toolbar'].addAction(self.export_txt)

    def export_bin_action(self):
        dialog = ExportDialog(self.mainwindow, ExportDialog.TIF)
        if dialog.exec_():
            result:ResultPointLayer = dialog.result_layer
            out = dialog.out_path
            if result:
                shutil.copy(result.result_path['path'], out)
                self.message_box.info('导出成功')

    def export_txt_action(self):
        dialog = ExportDialog(self.mainwindow, ExportDialog.TXT)
        if dialog.exec_():
            result = dialog.result_layer
            out = dialog.out_path
            if result:
                shutil.copy(result.path, out)
                self.message_box.info('导出成功')

    def export_rgb_action(self):
        dialog = ExportDialog(self.mainwindow, ExportDialog.TIF)
        btn = QPushButton('选择样式文件')
        style_path = ''
        def select_style():
            select_file = QFileDialog.getOpenFileName(self, '选择样式文件', '*.*')
            if select_file[0]:
                # style_path.setText(select_file[0])
                style_path = select_file[0]
                btn.setText(select_file[0])
        btn.clicked.connect(select_style)
        dialog.layout().addWidget(btn)

        default_style = [ 
            [0, 0, 255, 0],
            [1, 255, 0, 0]
        ]

        if style_path is '':
            style = default_style
        else:
            style = np.loadtxt(style_path, comments='#', delimiter=',')
            style = style.tolist()

        if dialog.exec_():
            result = dialog.result_layer.path
            out = dialog.out_path
            self.render(style, result, out)

    def render(self, style:list, path, out):
        data = gdal_array.LoadFile(path)
        if len(data.shape) == 3:
            data = data[..., 0]
        
        def get_color(v):
            first_color = []
            second_color = []
            first_value = 0
            second_value = 1
            for s in style:
                if s[0] <= v:
                    first_value = s[0]
                    first_color = s[1:]
                else:
                    second_value = s[0]
                    second_color = s[1:]
                    break
            if second_value == -1:
                return np.array(style[-1][1:])

            first_dis = (v - first_value) / (second_value - first_value)
            second_dis = (second_value - v) /  (second_value - first_value)
            first_color = np.array(first_color)
            second_color = np.array(second_color)
            color = first_color* first_dis + second_color * second_dis
            return color
    

        get_color = np.frompyfunc(get_color, nin=1, nout=1)

        Y, X = data.shape
        rgbs = get_color(data.reshape(-1,))
        rgbs = np.stack(rgbs).reshape((Y, X, 3)).astype(np.uint8)

        driver = gdal.GetDriverByName('GTiff')
        # out_tif = os.path.join(Project().other_path, 'temp.tif')
        out_ds = driver.Create(out, X, Y, 3, gdal.GDT_Byte)
        ds = gdal.Open(path)
        geo = ds.GetGeoTransform()
        proj = ds.GetProjectionRef()
        out_ds.SetGeoTransform(geo)
        out_ds.SetProjection(proj)

        for i in range(3):
            out_ds.GetRasterBand(i+1).WriteArray(rgbs[..., i])
        
        del out_ds