import cv2
import numpy as np
import mwTools

# Define the fields to process
fields = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]
    
# Load the coverageInfo datasets
print('Loading coverageInfo datasets...')
with np.load('coverageInfo.npz') as coverageInfoData:
    coverageCenterList = coverageInfoData['coverageCenterList']

# Define the video codec and create VideoWriter object
print('Creating video...')
fourcc = cv2.VideoWriter_fourcc(*'H264')
frameRate = 60
frameWidth = 1920
frameHeight = 1080
videoOutput = cv2.VideoWriter('mw.avi', fourcc, frameRate, (frameWidth, frameHeight))

for field in fields:
    # Define the fields in both sides
    previousField = field - 1 if field - 1 >= 0 else 14
    nextField = field + 1 if field + 1 < 15 else 0
    
    # Load the bgr files
    print('Loading bgr files...')
    bgrData = None
    
    for i in [previousField, field, nextField]:
        with np.load('fied'+str(i)+'.npz') as data:
            if bgrData is None: 
                bgrData = data['bgrData']
            else:
                bgrData = np.hstack((bgrData, data['bgrData']))
   
    # Calculate the smoothed coverage center for each column
    print('Obtaining the smoothed coverage center...')
    coverageCenter = np.hstack((coverageCenterList[previousField], coverageCenterList[field], coverageCenterList[nextField]))
    coverageCenter = mwTools.smooth(coverageCenter, 5000)
    
    # Resize the datasets
    print('Resizing the datasets...')
    scaleFactor = 0.6
    bgrData = cv2.resize(bgrData, None, fx=scaleFactor, fy=scaleFactor, interpolation=cv2.INTER_CUBIC)
    coverageCenter = scaleFactor * np.array([coverageCenter[round(i/scaleFactor)] for i in range(bgrData.shape[1])])
    
    # Display and save the results
    fieldWidth = round(bgrData.shape[1]/3)
    fieldHeight = bgrData.shape[0]
    startCol = fieldWidth - round(frameWidth/2)
    endCol = startCol + fieldWidth
    colStep = 2
    colPercentageStep = round(fieldWidth/20) * colStep
    frame = np.copy(bgrData[0:frameHeight, 0:frameWidth])
    exit = False
    
    for col in range(startCol, endCol, colStep):    
        # Copy the region of interest to the frame dataset
        row = round(coverageCenter[col + round(frameWidth/2)] - frameHeight/2)
        np.copyto(frame, bgrData[row:row + frameHeight, col:col + frameWidth])
        
        # Obtain teh frame central coordinates
        fieldCol = col - startCol
        longitude = 180.0 - (24.0*field + (fieldCol/scaleFactor)*3.2/3600)
        latitude = -((row + frameHeight/2 - fieldHeight/2)/scaleFactor)*3.2/3600
        mwTools.writeCoordinates(frame, round(0.82*frameWidth), round(0.05*frameHeight), longitude, latitude)
        
        # Display the frame
        cv2.imshow('PACS Milky Way', frame)
        
        # User interaction
        if cv2.waitKey(10) & 0xFF == ord('q'):
            exit = True
            break
        
        # Write the frame
        videoOutput.write(frame)
        
        # Write some status information
        if fieldCol % colPercentageStep == 0:
            print(round(fieldCol*100/fieldWidth), '%', sep='', end=' ... ', flush=True)
    
    print('Finished')
    del(bgrData)
    
    if exit: 
        break

videoOutput.release()
cv2.destroyAllWindows()