"""Component IMF, unresolved band fluxes and imposed spatial snapshots, not N-body ages."""
import numpy as np


def integral_power(low, high, slope):
    if abs(slope-1) < 1e-10:
        return np.log(high/low)
    return (high**(1-slope)-low**(1-slope))/(1-slope)


def inverse_power(uniform, low, high, slope):
    if abs(slope-1) < 1e-10:
        return low*(high/low)**uniform
    return (low**(1-slope)+uniform*(high**(1-slope)-low**(1-slope)))**(1/(1-slope))


def sample_imf(uniform, high_slope=2.3, low=.1, high=8):
    split=.5; factor=split**(high_slope-1.3)
    first=integral_power(low, split, 1.3)
    second=factor*integral_power(split, high, high_slope)
    fraction=first/(first+second)
    uniform=np.asarray(uniform)
    result=np.empty_like(uniform)
    left=uniform < fraction
    result[left]=inverse_power(uniform[left]/fraction, low, split, 1.3)
    result[~left]=inverse_power((uniform[~left]-fraction)/(1-fraction), split, high, high_slope)
    return result


def combine_magnitudes(first, second):
    return -2.5/np.log(10)*np.logaddexp(-.4*np.log(10)*np.asarray(first), -.4*np.log(10)*np.asarray(second))


def sample_plummer(uniform, cos_theta, phi, scale=1):
    radius=scale/np.sqrt(np.asarray(uniform)**(-2/3)-1)
    plane=np.sqrt(1-np.asarray(cos_theta)**2)
    return radius[:,None]*np.column_stack([plane*np.cos(phi), plane*np.sin(phi), cos_theta])


def pair_components(masses, system_binary_fraction, rng):
    if not 0 <= system_binary_fraction <= 1:
        raise ValueError('Binary system fraction outside [0,1]')
    number_pairs=int(np.floor(system_binary_fraction*len(masses)/(1+system_binary_fraction)))
    indices=rng.permutation(len(masses))
    systems=[indices[2*i:2*i+2] for i in range(number_pairs)]
    systems.extend(indices[i:i+1] for i in range(2*number_pairs, len(masses)))
    return systems


def synthesize_population(model, seed=32, binary_fraction=.35, high_slope=2.3, segregation=.4):
    if not 0<=segregation<=1:
        raise ValueError('Imposed segregation outside [0,1]')
    rng=np.random.default_rng(seed)
    masses=sample_imf(rng.random(1024),high_slope)
    systems=pair_components(masses,binary_fraction,rng)
    # Fixed per-component spatial latents keep the same centre for a surviving
    # system anchor when multiplicity changes; do not reshuffle every direction.
    vectors=sample_plummer(rng.uniform(1e-10,1-1e-10,len(masses)),rng.uniform(-1,1,len(masses)),rng.uniform(0,2*np.pi,len(masses)))
    records=[]
    for index,components in enumerate(systems):
        initial_mass=float(masses[components].sum())
        position=vectors[int(components[0])]*2*(initial_mass/.5)**(-segregation)
        distance=float(np.linalg.norm(position+np.array([0,0,850.])))
        if position[2]+850<=0:raise ValueError('Synthetic system is behind the observer')
        angular_radius=float(np.degrees(np.arctan2(np.hypot(position[0],position[1]),position[2]+850)))
        row={'componentIndices':components.tolist(),'initialComponentMassesMsun':masses[components].tolist(),
             'initialSystemMassMsun':initial_mass,'positionPc':position.tolist(),'distancePc':distance,
             'angularRadiusDegrees':angular_radius,'inAperture':angular_radius<.3,'fluxSupported':True}
        magnitudes=[]
        for mass in masses[components]:
            if not model['Mini'][0]<=mass<=model['Mini'][-1]:row['fluxSupported']=False;break
            lower=np.searchsorted(model['Mini'],mass,side='right')-1
            upper=min(lower+1,len(model['Mini'])-1)
            if model['label'][lower]>5 or model['label'][upper]>5 or model['label'][lower]!=model['label'][upper]:
                row['fluxSupported']=False;break
            magnitudes.append([float(np.interp(mass,model['Mini'],model[band])) for band in ['Gmag','G_BPmag','G_RPmag']])
        if row['fluxSupported']:
            flux=np.sum(10**(-.4*np.array(magnitudes)),axis=0)
            combined=-2.5*np.log10(flux);apparent=combined+5*np.log10(distance/10)
            row.update({'componentAbsoluteMagnitudes':magnitudes,'summedBandFluxOverZeroPoints':flux.tolist(),
                        'absoluteMagnitudes':combined.tolist(),'apparentMagnitudes':apparent.tolist(),
                        'BPminusRP':float(apparent[1]-apparent[2]),'detected':bool(row['inAperture'] and apparent[0]<18.5)})
        else:row['reason']='At least one component is outside luminous support or spans a phase boundary; no zero-flux remnant or bridge invented'
        records.append(row)
    return {'kind':'own single-age synthetic component population, imposed spatial snapshot',
            'componentInitialMassesMsun':masses.tolist(),'binarySystemFractionRequested':binary_fraction,
            'binarySystemFractionRealized':sum(len(row['componentIndices'])==2 for row in records)/len(records),
            'highMassIMFSlope':high_slope,'imposedSegregation':segregation,'systems':records,
            'distanceCentrePc':850,'scaleReferencePc':2,'apertureDegrees':.3,
            'dynamicalEvolutionComputed':False,'remnantLuminositiesComputed':False,
            'units':'Solar birth masses, pc positions/distances; Vega magnitudes and dimensionless per-band flux/zero-point ratios',
            'boundary':'All components are sampled before pairing, preserving the component IMF and initial mass. Unresolved systems share a centre. No physical binary orbit, multiplicity evolution or N-body clock is implied.'}