import tkinter
from  astropy.io import fits
import numpy as np
import pickle
from  astropy.io import fits
import matplotlib.pyplot as plt
import os
from matplotlib import widgets
def fitsData(fitsName):
    with fits.open(fitsName) as hdul:
        hdul.info()
        data = hdul[0].data
        shape = data.shape
        data = data.reshape((shape[-1], shape[-2]))
        return data
def plotContour(data, ax,base):
    """将一个矩阵的等高线画在对应的ax上,用base来衡量等高线的level"""
    shape = data.shape
    y, x = np.mgrid[0:shape[0],
           0:shape[1]]
    ax.contour(x, y, data, [-base, base, base * 1.41, base * 2, base * 2.83, base * 4,
                            base * 5.65, base * 7.99, base * 11.3, base * 16, base * 22.6, base * 32, base * 45.2,
                            base * 64], colors='black')
def plotBox(x,y,width,height,ax):
    """定义画方框的函数,依次对应方框左下角横纵坐标，宽度，高度以及对应的ax"""
    ax.hlines(y=[y, y+height]
               , xmin=x
               , xmax=x+width
              )
    ax.vlines(x=[x,x+width]
                    , ymin=y
                    , ymax=y+height
              )

def crossCorrelation(data1,data2,outFileName):
    """
    :param data1: 第一个二维矩阵数据
    :param data2: 第二个二维矩阵数据
    :param outFileName: 输出pkl文件的名字，要写后缀
    :param data1On: 默认为数据1在上，设置为False的话数据2在上
    :return: 两个数据的cross-correlation，并将每个点的相关数据存储在一个pkl文件里
    """
    shape1=data1.shape
    shape2=data2.shape
    """取数据1的0行0列为初始点，让这个点在数据2的0行0列开始逐行逐列平移
    现在初始点移动到了数据2的第x行第y列
    """
    rho=[]
    for i in range(shape2[0]):
        for j in range(shape2[1]):
            """
            包含的数据1的数据为数据1的第0行到shape1[0]-x行，第0列到shape1[1]-y列
            包含数据2的数据为数据2的第x行到最后一行，第y列到最后一列.
            """
            subData1=data1[0:shape1[0]-i,0:shape1[1]-j]
            subData2=data2[i:,j:]
            subShape1=subData1.shape
            subShape2=subData2.shape
            subMean1=subData1.mean()
            subMean2=subData2.mean()
            sum1=0
            sum2=0
            sum3=0
            for x in range(subShape1[0]):
                for y in range(subShape1[1]):
                    sum1=sum1+(subData1[x,y]-subMean1)*(subData2[-x,-y]-subMean2)
                    sum2=sum2+(subData1[x,y]-subMean1)**2*(subData2[-x,-y]-subMean2)**2
            rho.append((i,j,sum1/(sum2**0.5)))
            print(i,j)
    with open(outFileName,'wb') as file:
        pickle.dump(rho,file)
    print("Cross_Correlation++++++++++++++++done")
class FitsGUI():
    def __init__(self):
        """导入两个fits图片的数据"""
        fitsName1='8GHz.fits'
        fitsName2='22GHz.fits'
        self.data1=fitsData(fitsName1)
        self.data2 = fitsData(fitsName2)
        self.shape1=self.data1.shape
        self.shape2 = self.data1.shape
        ###############################################################################
        """
        figure大小设置为10*10，分为四块，下面两块用来显示数据1与数据2，对应的ax为4*4，左侧与下侧空隙均为0.5.
        一开始图1与2显示自己对应的fits,ax3用来显示最终相关的结果
        """
        self.fig=plt.figure(figsize=(10,10))
        self.ax1 = self.fig.add_axes([0.5 / 10, 0.5 / 10, 4/10, 4/10])
        self.ax2 = self.fig.add_axes([5.5 / 10, 0.5 / 10, 4 / 10, 4 / 10])
        self.ax3 = self.fig.add_axes([5.5 / 10, 5.5 / 10, 4 / 10, 4 / 10])
        plotContour(self.data1,self.ax1,base=0.0015)
        plotContour(self.data2, self.ax2,base=0.001)

        ###############################################################################
        """"
        滑块长度均为4，宽度为0.15，用来移动方框的滑块的长度设置为图片的长度
        """
        ##############################################
        """图1滑块设置，用来控制图1中方框左下角点的坐标"""
        self.ax1XSliderAxes = self.fig.add_axes([0.5/10, 4.75/10, 4/10, 0.15/10])
        self.ax1XSlider = widgets.Slider(ax=self.ax1XSliderAxes, valmin=0, valmax=self.shape1[1], valstep=1, label='X')
        self.ax1YSliderAxes = self.fig.add_axes([4.75 / 10, 0.5 / 10, 0.15 / 10, 4 / 10])
        self.ax1YSlider = widgets.Slider(ax=self.ax1YSliderAxes, valmin=0, valmax=self.shape1[0], valstep=1, label='Y'
                                         ,orientation='vertical')
        ##############################################
        """图2滑块设置，用来控制图2中方框左下角点的坐标"""
        self.ax2XSliderAxes = self.fig.add_axes([5.5 / 10, 4.75 / 10, 4 / 10, 0.15 / 10])
        self.ax2XSlider = widgets.Slider(ax=self.ax2XSliderAxes, valmin=0, valmax=self.shape2[1], valstep=1, label='X')
        self.ax2YSliderAxes = self.fig.add_axes([9.75 / 10, 0.5 / 10, 0.15 / 10, 4 / 10])
        self.ax2YSlider = widgets.Slider(ax=self.ax2YSliderAxes, valmin=0, valmax=self.shape2[0], valstep=1, label='Y'
                                         , orientation='vertical')
        #############################################
        """设置用来控制方框大小的滑块，这两个滑块放在图片左上方，分别代表方框的宽和高,上面控制方框高度，下面控制宽度，宽与高的最大值与图片长度相同"""
        self.BoxWidthAxes=self.fig.add_axes([0.5/10, 9.2 / 10, 4 / 10, 0.15 / 10])
        self.BoxWidthSlider=widgets.Slider(ax=self.BoxWidthAxes, valmin=0, valmax=self.shape2[1], valstep=1
                                               , label='Width')

        self.BoxHeightAxes = self.fig.add_axes([0.5 / 10, 9.65 / 10, 4 / 10, 0.15 / 10])
        self.BoxHeightSlider = widgets.Slider(ax=self.BoxHeightAxes, valmin=0, valmax=self.shape2[0], valstep=1
                                                 ,label='Height')
        #############################################
        """所有滑块的变化都对应函数sliderChange"""
        self.ax1XSlider.on_changed(self.sliderChange)
        self.ax1YSlider.on_changed(self.sliderChange)
        self.ax2XSlider.on_changed(self.sliderChange)
        self.ax2YSlider.on_changed(self.sliderChange)
        self.BoxHeightSlider.on_changed(self.sliderChange)
        self.BoxWidthSlider.on_changed(self.sliderChange)
        ###############################################################################
        """选定好区域后，点击开始互相关按钮，开始计算最终成图"""
        self.startbuttonAxes=self.fig.add_axes([2/10, 8 / 10, 1 / 10, 0.5 / 10])
        self.startButton=widgets.Button(ax=self.startbuttonAxes,label='START')
        self.startButton.on_clicked(self.buttonClick)


        ###############################################################################
        plt.show()
    def sliderChange(self,val):
        """开始画方框以及显示fits图片,也就是设置各个滑块的交互函数"""
        self.ax1.cla()
        self.ax2.cla()
        plotContour(self.data1,self.ax1,base=0.0015)
        plotContour(self.data2, self.ax2,base=0.001)

        plotBox(ax=self.ax1,x=self.ax1XSlider.val,y=self.ax1YSlider.val,height=self.BoxHeightSlider.val
                ,width=self.BoxWidthSlider.val)
        plotBox(ax=self.ax2, x=self.ax2XSlider.val, y=self.ax2YSlider.val, height=self.BoxHeightSlider.val
                ,width=self.BoxWidthSlider.val)

        # print('数据1从第{}到第{}行，第{}列到第{}列'.format(self.ax1YSlider.val
        #                                       ,self.ax1YSlider.val+self.BoxHeightSlider.val
        #                                       ,self.ax1XSlider.val
        #                                       ,self.ax1XSlider.val+self.BoxWidthSlider.val))
        #
        # print('数据2从第{}到第{}行，第{}列到第{}列'.format(self.ax2YSlider.val
        #                                       , self.ax2YSlider.val + self.BoxHeightSlider.val
        #                                       , self.ax2XSlider.val
        #                                       , self.ax2XSlider.val + self.BoxWidthSlider.val))
        self.subData1=self.data1[self.ax1YSlider.val:self.ax1YSlider.val+self.BoxHeightSlider.val
                 ,self.ax1XSlider.val:self.ax1XSlider.val+self.BoxWidthSlider.val]
        self.subData2=self.data2[self.ax2YSlider.val:self.ax2YSlider.val + self.BoxHeightSlider.val
                 ,self.ax2XSlider.val:self.ax2XSlider.val + self.BoxWidthSlider.val]


    def buttonClick(self,click):
        crossCorrelation(self.subData1,self.subData2,outFileName="corss.pkl")
        with open('corss.pkl', 'rb') as file1:
            data1 = pickle.load(file1)
            data1 = sorted(data1, key=lambda x: x[-1], reverse=True)
            print(data1[0])
if __name__=='__main__':
    FitsGUI()