#!/usr/bin/env python3

import re
import math
import shutil
import subprocess
import tempfile
from pathlib import Path

from gerbonara import LayerStack
from gerbonara.utils import setup_svg, Tag, MM, rotate_point
import gerbonara.cad.kicad.pcb as pcb
from PIL import Image
import click
import shapely
import numpy as np
from scipy.optimize import minimize

# Default density according to https://www.avx.com/docs/techinfo/CeramicCapacitors/yngsmodu.pdf
density_x7r = 5.8 # g/cm^3
# Default density according to https://en.wikipedia.org/wiki/FR-4
density_fr4 = 1.850 # g/cm^3

density_copper = 8.96 # g/cm^3

def volume_mass(w:'mm', l:'mm', density:'g/cm^3', thickness:'mm'=None):
    """ Calculate approximate mass of a component of a given size. When no thickness is given, thickness = width is
    assumed. Lengths are in mm, density in g/cm^3 and the result will be in mg. """
    thickness = thickness or w
    return w*l*thickness * 1e-3 * density * 1e3

# masses for some common components in mg
# approximate! Especially for capacitors, there can be significant differences between component height (and thus mass)
# for parts with the same footprint.
common_masses = {
        'C_0201_0603Metric': volume_mass(0.3, 0.6, density_x7r),
        'C_0402_1005Metric': volume_mass(0.5, 1.0, density_x7r),
        'C_0603_1608Metric': volume_mass(0.8, 1.5, density_x7r),
        'C_0805_2012Metric': volume_mass(1.3, 2.0, density_x7r),
        'C_1206_3216Metric': volume_mass(1.5, 3.0, density_x7r),
        'C_1210_3225Metric': volume_mass(2.5, 3.2, density_x7r),
        # Resistor masses from https://www.vishay.com/docs/20035/dcrcwe3.pdf
        'R_0402_1005Metric': 0.65,
        'R_0603_1608Metric': 2.0,
        'R_0805_2012Metric': 5.5,
        'R_1206_3216Metric': 10,
        }


def angle_difference(a, b):
    return (b - a + math.pi) % (2*math.pi) - math.pi

def parse_mass(s):
    if not (m := re.match(r'^([0-9.,]*)([mkuµ]?)[g]?', s)):
        raise ValueError(f'Invalid mass "{s}"')
    si_prefix = {'m': 1e-3, 'u': 1e-6, 'µ': 1e-6, 'k': 1e3}
    return float(m.group(1)) * si_prefix.get(m.group(2), 1)


def get_pixel_size(img, phys_w, phys_h):
    px_w, px_h = img.size
    pixel_size_x, pixel_size_y = phys_w/px_w, phys_h/px_h
    assert math.isclose(pixel_size_x, pixel_size_y, rel_tol=1e-2, abs_tol=0.05)
    return (pixel_size_x + pixel_size_y) / 2


def pixel_inertia(svg, phys_w, phys_h, phys_x0, phys_y0, area_density):
    with tempfile.NamedTemporaryFile(suffix='.png') as tmp_png:
        tmp_png = Path(tmp_png.name)

        subprocess.run(['resvg', '--background', 'white', str(svg), str(tmp_png)], check=True)
        with Image.open(tmp_png) as img:
            pixel_size = get_pixel_size(img, phys_w, phys_h) # [m]
            px_w, px_h = img.size
            img = 1 - np.array(img.convert('L'), dtype=float) / 255
            pixel_weight = pixel_size**2 * area_density # [g]
            img *= pixel_weight # [g]

            tens_y = np.zeros(shape=(px_h, 3, 3))
            tens_x = np.zeros(shape=(px_w, 3, 3))
            mass_y = np.zeros(px_h)
            mass_x = np.zeros(px_w)
            com_y = np.zeros(shape=(px_h, 2))
            com_x = np.zeros(shape=(px_w, 2))

            for y in range(px_h):
                for x in range(px_w):
                    m = img[y, x]
                    x_phys = (x + 0.5) * pixel_size - phys_x0 # [m]
                    y_phys = (y + 0.5) * pixel_size - phys_y0 # [m]

                    tens_x[x][0,0] = m * y_phys**2
                    tens_x[x][1,1] = m * x_phys**2
                    tens_x[x][2,2] = m * x_phys**2 + y_phys**2
                    tens_x[x][0,1] = tens_x[x][1,0] = m * x_phys * y_phys

                    mass_x[x] = m

                    com_x[x] = (x_phys*m, y_phys*m)

                tens_y[y] = tens_x.sum(axis=0)
                mass_y[y] = mass_x.sum()
                com_y[y] = com_x.sum(axis=0)

            tens_solder = tens_y.sum(axis=0)
            mass_solder = mass_y.sum() 
            center_of_mass = com_y.sum(axis=0) / mass_solder
            return mass_solder, center_of_mass, tens_solder


@click.command()
@click.argument('kicad_pcb', type=click.Path(file_okay=True, dir_okay=False, path_type=Path))
@click.argument('gerber', type=click.Path(file_okay=False, dir_okay=True, path_type=Path))
@click.argument('pcb_out', required=False, type=click.Path(file_okay=True, dir_okay=False, path_type=Path))
@click.option('--board-thickness', type=float, default=0.75, help='PCB thickness in mm')
@click.option('--fr4-density', type=float, default=density_fr4, help='FR-4 density in g/cm^3')
@click.option('--stencil-thickness', type=float, default=0.125, help='Stencil thickness in mm')
@click.option('--copper-thickness', type=float, default=35, help='Stencil thickness in µm')
# Default solder paste density from:
# - https://www.kester.com/Portals/0/Documents/Knowledge%20Base/Calculating%20Solder%20Paste%20Volume%20Percent.pdf
# - http://www.chipquik.com/datasheets/SMD291SNL10.pdf
@click.option('--solder-paste-density', type=float, default=4.2, help='Solder paste density in g/cm^3')
@click.option('--solder-paste-metal-content', type=float, default=87.0, help='Solder paste metal content by weight in percent. The default value is for Chipquik SMD291SNL lead-free SAC305 solder paste.')
def analyze_board(kicad_pcb, gerber, pcb_out, board_thickness, fr4_density, stencil_thickness, copper_thickness, solder_paste_density, solder_paste_metal_content):
    board = pcb.Board.open(kicad_pcb)
    stack = LayerStack.open(gerber)
    print()

    (min_x, min_y), (max_x, max_y) = bounds = stack.outline.bounding_box(MM, default=((0, 0), (0, 0)))
    phys_w, phys_h = (max_x-min_x)/1e3, (max_y-min_y)/1e3 # [m]
    # Gerber flips the image Y axis compared to pillow and other modern computer graphics
    phys_x0, phys_y0 = -min_x/1e3, max_y/1e3 # [m]
    stencil_thickness /= 1e3 # [m]
    solder_paste_density *= 1e6 / 1e3 # [kg/m^3]
    solder_paste_metal_content /= 100 # [1]
    board_thickness /= 1e3 # [m]
    copper_thickness /= 1e6 # [m]
    fr4_density *= 1e6 / 1e3 # [kg/m^3]

    with tempfile.NamedTemporaryFile(suffix='.svg') as tmp_svg:
        tmp_svg = Path(tmp_svg.name)
        tmp_svg.write_text(str(stack['top paste'].to_svg(force_bounds=bounds)))

        area_density = stencil_thickness * solder_paste_density * solder_paste_metal_content # [kg/m^2]
        solder_mass, solder_center_of_mass, solder_tens = pixel_inertia(tmp_svg, phys_w, phys_h, phys_x0, phys_y0, area_density)

        print(f'Solder mass: {solder_mass*1e6:.3f} mg')
        print(f'Solder inertia:\n', solder_tens)
        print(f'Solder center of mass: {solder_center_of_mass[0]*1e3:.3f}, {solder_center_of_mass[1]*1e3:.3f} mm')
        print()

    with tempfile.NamedTemporaryFile(suffix='.svg') as tmp_svg:
        tmp_svg = Path(tmp_svg.name)

        volume_density = density_copper * 1e6 / 1e3 # [kg/m^3]
        area_density = copper_thickness * volume_density # [kg/m^2]

        tmp_svg.write_text(str(stack['top copper'].to_svg(force_bounds=bounds)))
        top_copper_mass, top_copper_center_of_mass, top_copper_tens = pixel_inertia(tmp_svg, phys_w, phys_h, phys_x0, phys_y0, area_density)

        tmp_svg.write_text(str(stack['bottom copper'].to_svg(force_bounds=bounds)))
        bottom_copper_mass, bottom_copper_center_of_mass, bottom_copper_tens = pixel_inertia(tmp_svg, phys_w, phys_h, phys_x0, phys_y0, area_density)

        copper_mass = top_copper_mass + bottom_copper_mass
        copper_center_of_mass = (top_copper_center_of_mass * top_copper_mass +
                                 bottom_copper_center_of_mass * bottom_copper_mass) / copper_mass
        copper_tens = top_copper_tens + bottom_copper_tens

        print(f'Copper mass: {copper_mass*1e3:.3f} g')
        print(f'Copper inertia:\n', copper_tens)
        print(f'Copper center of mass: {copper_center_of_mass[0]*1e3:.3f}, {copper_center_of_mass[1]*1e3:.3f} mm')
        print()

    with tempfile.NamedTemporaryFile(suffix='.svg') as tmp_svg:
        tmp_svg = Path(tmp_svg.name)
        path = Tag('path', d=stack.outline_svg_d(), fill='black')
        tmp_svg.write_text(str(setup_svg([path], bounds)))

        area_density = board_thickness * fr4_density # [kg/m^2]
        fr4_mass, fr4_center_of_mass, fr4_tens = pixel_inertia(tmp_svg, phys_w, phys_h, phys_x0, phys_y0, area_density)

        print(f'FR-4 mass: {fr4_mass*1e3:.3f} g')
        print(f'FR-4 inertia:\n', fr4_tens)
        print(f'FR-4 center of mass: {fr4_center_of_mass[0]*1e3:.3f}, {fr4_center_of_mass[1]*1e3:.3f} mm')
        print()

    component_mass = np.zeros(len(board.footprints))
    component_tens = np.zeros(shape=(len(board.footprints), 3, 3))
    component_com  = np.zeros(shape=(len(board.footprints), 2))
    for i, fp in enumerate(board.footprints):
        if (m := fp.property_value('mass', None)):
            m = parse_mass(m) / 1e3

        else:
            for key, value in common_masses.items():
                if key in fp.name:
                    m = value/1e6 # [kg]
                    break
            else:
                m = 0

        x, y = fp.at.x, fp.at.y # [mm]
        x -= board.setup.aux_axis_origin.x
        y -= board.setup.aux_axis_origin.y
        x, y = x*1e-3, y*1e-3 # [m]

        component_tens[i][0,0] = m * y**2
        component_tens[i][1,1] = m * x**2
        component_tens[i][2,2] = m * x**2 + y**2
        component_tens[i][0,1] = component_tens[i][1,0] = m * x * y

        component_mass[i] = m

        component_com[i]  = (m*x, m*y)

    component_mass = component_mass.sum()
    component_tens = component_tens.sum(axis=0)
    component_center_of_mass = component_com.sum(axis=0) / component_mass
    
    print(f'Component mass: {component_mass*1e3:.3f} g')
    print(f'Component inertia:\n', component_tens)
    print(f'Component center of mass: {component_center_of_mass[0]*1e3:.3f}, {component_center_of_mass[1]*1e3:.3f} mm')
    print()

    total_mass = fr4_mass + solder_mass + component_mass + copper_mass
    total_tens = fr4_tens + solder_tens + component_tens + copper_tens
    total_center_of_mass = (
            fr4_center_of_mass * fr4_mass +
            solder_center_of_mass * solder_mass +
            component_center_of_mass * component_mass +
            copper_center_of_mass * copper_mass) / total_mass
    imbalance_angle = math.atan2(*total_center_of_mass[::-1])
    print(f'Total mass: {total_mass*1e3:.3f} g')
    print(f'Total inertia:\n', total_tens)
    print(f'Total center of mass: {total_center_of_mass[0]*1e3:.3f}, {total_center_of_mass[1]*1e3:.3f} mm')
    print(f'Imbalance angle: {math.degrees(imbalance_angle):.2f} °')
    print(f'Imbalance radius: {math.hypot(*total_center_of_mass)*1e3:.2f} mm')
    print()

    if not pcb_out: # No auto-balancing requested
        return

    volume_density = density_copper * 1e6 / 1e3 # [kg/m^3]
    area_density = copper_thickness * volume_density # [kg/m^2]

    balancing_pads = []
    for fp in board.footprints:
        if not fp.name.startswith('balancing-footprints:bal-'):
            continue
        fp_x, fp_y = fp.at.x - board.setup.aux_axis_origin.x, fp.at.y - board.setup.aux_axis_origin.y
        fp_x, fp_y = fp_x*1e-3, fp_y*1e-3 # [m]
        fp_rot = math.radians(fp.at.rotation or 0)

        for pad in fp.pads:
            if pad.shape not in (pcb.Atom.rect, pcb.Atom.circle):
                print(f'Pad {pad.numebr} of balancing footprint {fp.get_property("Reference")} has unsupported pad shape {pad.shape}. Ignoring.')
                continue

            x, y = pad.at.x, pad.at.y
            x, y = x*1e-3, y*1e-3 # [m]
            x, y = rotate_point(x, y, -fp_rot)
            x, y = x+fp_x, y+fp_y

            a = math.atan2(y, x)
            a_diff = angle_difference(imbalance_angle, a)


            w, h = pad.size.x*1e-3, pad.size.y*1e-3 # [m]
            if pad.shape == pcb.Atom.rect:
                m = w*h*area_density
            else:
                m = max(w, h)**2 / 4 * math.pi * area_density

            balancing_pads.append((pad, a, a_diff, x, y, w, h, m))

    if not balancing_pads:
        print('Error: Cannot balance board: No balancing footprint instances found.')
        return

    com_wo_bal_fps = total_center_of_mass * total_mass
    m_tot = total_mass
    for i, (pad, a, a_diff, x, y, w, h, m) in enumerate(balancing_pads):
        #print(f'Balancing footprint {i} at {x*1e3:.1f}, {y*1e3:.1f} mm')
        #print(f'  angle: {math.degrees(a):.2f} °')
        #print(f'  angle difference: {math.degrees(a_diff):.2f} °')
        #print(f'  radius: {math.hypot(x, y)*1e3:.1f} mm')
        #print(f'  weight: {m*1e6:.0f} mg')
        #print()

        com_wo_bal_fps -= np.array([x, y]) * m
        m_tot -= m
    com_wo_bal_fps /= m_tot

    def objective_function(x):
        xs = x
        total_com, m_bal = com_wo_bal_fps*m_tot, 0
        for i, (pad, a, a_diff, x, y, w, h, m) in enumerate(balancing_pads):
            m *= xs[i]
            total_com += np.array([x, y])*m
            m_bal += m
        total_com /= (m_tot + m_bal)

        r = math.hypot(*total_com)
        #mean_change = 1 - xs.mean()

        return r #+ mean_change * 1e-3

    x0 = np.ones(len(balancing_pads))
    res = minimize(objective_function, x0, bounds=[(0.005, 1.0)]*len(x0), tol=1e-12)
    print(res)
    print(f'Residual imbalance radius: {objective_function(res.x)*1e3:.2f} mm')
    print('Adjusting balancing footprints:')
    for i, (pad, a, a_diff, x, y, w, h, m) in enumerate(balancing_pads):
        f = math.sqrt(res.x[i])
        pad.size.x *= f
        pad.size.y *= f
        print(f'[{i:> 3}] by {100-res.x[i]*100:> 6.2f} %', end=('\n' if i%4 == 3 else '  '))

    board.write(pcb_out)


if __name__ == '__main__':
    analyze_board()

