from PIL import Image, ImageDraw, ImageOps
from scipy.optimize import curve_fit
import rawpy
import matplotlib.pyplot as plt
import numpy as np
import math

def findlines (im,doplot=False):
    rows = list()
    lines = {'red':{'nm':611.43,'lums':[],'cols':[]},
         'green':{'nm':545.98,'lums':[],'cols':[]},
         'blue':{'nm':435.85,'lums':[],'cols':[]}}
    lum = {}
    col = {}
    for r in range(12,im.height,30):
        # Find spectral lines
        # Extract three rows and resize them into one to reduce noise
        ims = im.crop((0,r-1,im.width,r+2)).resize((im.width,1))
        # Now extract the luminance of the combined three rows
        iml = ims.convert('L')
        imlv = list(iml.getdata())

        # Find the peaks in the luminance
        lum['green'] = max(imlv[im.width//3:im.width*2//3])
        if lum['green']<32 or lum['green']>192: continue # Skip rows with out-of-range luminance

        col['green'] = imlv.index(lum['green'],im.width//3,im.width*2//3)
        endofgreen = col['green']+100
        lum['red'] = max(imlv[endofgreen:])
        col['red'] = imlv.index(lum['red'],endofgreen)
        col2nm = (lines['red']['nm']-lines['green']['nm'])/(col['red']-col['green'])
        predblue = int(col['green']-(lines['green']['nm']-lines['blue']['nm'])/col2nm)
        lum['blue'] = max(imlv[predblue-100:predblue+100])
        col['blue'] = imlv.index(lum['blue'],predblue-100,predblue+100)

        rows.append(r)
        for color in lines:
            lines[color]['lums'].append(lum[color])
            lines[color]['cols'].append(col[color])

        if doplot:
            fig,ax = plt.subplots()
            ax.plot(range(len(imlv)),imlv)
            ax.plot(list(col[color] for color in lines),list(lum[color] for color in lines),'x')
    return rows,lines

fn = 'C:/Users/Hobbyraum/Documents/Spektrometer/esl8w.v4.png'
im = Image.open(fn)
iw = im.width
ih = im.height
rows,lines = findlines(im,False)

# Calculate matrix fit of nm per row (polygonal) and col (linear)

def rc2nm(row,col,a,b,c,d): 
    return a*row*row+b*row+c*col+d

def fitme(rc,a,b,c,d): 
    row,col = rc.T
    return a*row*row+b*row+c*col+d

xdata = list()
ydata = list()
for lc in lines.values():
    xdata.extend(zip(rows,lc['cols']))
    ydata.extend(list(lc['nm'] for row in rows))
popt,_ = curve_fit(fitme,np.array(xdata),np.array(ydata))

nmrc = dict()
for row in range(ih//4,ih*3//4):
    for col in range(iw):
        nm = int(rc2nm(row,col,*popt))
        nmrc.setdefault(nm,[]).append((row,col))

fn = 'halogen70w.v4.g2.dng'
rp = rawpy.imread(fn)
ri = rp.raw_image
rc = rp.raw_colors

cnmri = dict()
for nm in sorted(nmrc):
    for row,col in nmrc[nm]:
        c = rc[row][col]
        l = ri[row][col]
        cnmri.setdefault(c,{}).setdefault(nm,[]).append(l)

colors = rp.color_desc.decode('ascii')
colorspecs = {0:'#ff0000',1:'#88cc00',2:'#0000ff',3:'#00cc88'}
fig,ax = plt.subplots()
for c in sorted(cnmri):
    nmri = cnmri[c]
    cl = f'Bayer filter color {colors[c]} ({c:d})'
    cs = colorspecs[c]
    ax.scatter(nmri.keys(),list(sum(nmri[nm])/len(nmri[nm]) for nm in nmri),s=list(math.log(len(nmri[nm])) for nm in nmri),color=cs,label=cl)
ax.set_xlabel('nm')
ax.set_ylabel('avg raw luminosity')
ax.legend()
ax.set_title(fn)
plt.show()