import mido
import copy
import sys
from mido import MidiFile, MidiTrack, merge_tracks
from collections import defaultdict
from dataclasses import dataclass, field

@dataclass
class NoteEvent:
    start: int
    end: int
    note: int
    velocity: int
    orgchan: int#vioce数据开头有，但是压实后会丢失，且需要在压实后根据此数据复制OrgChanControl属性
    timesec: float#将start tick转成秒数存储

@dataclass
class OrgChanControl:
    tick: int
    program: int | None = None
    pitchwheel: int | None = None
    controls: list = field(default_factory=list)
    aftertouch: list = field(default_factory=list)
    
@dataclass
class Voice:
    events: list[NoteEvent]
    
@dataclass
class PackedNote:
    start: int
    end: int
    note: int
    velocity: int
    outchan: int      # 压实后的通道
    orgchan: int      # 原始通道（用于回填控制）

def to_midi0(mid: MidiFile) -> MidiFile:
    print ("1. 转为单轨Midi ",end="")
    if len(mid.tracks) == 1:
        print ("不用转换 OK")
        return mid

    out = MidiFile(ticks_per_beat=mid.ticks_per_beat)
    track = MidiTrack()
    out.tracks.append(track)

    for msg in merge_tracks(mid.tracks):
        track.append(msg)
    print ("已转换 OK")
    return out

def split_channel_polyphony(track,tpb):
    print ("2. 复音展开 ",end="")
    # ch -> list[voice]
    voices = defaultdict(list)

    # (ch, note) -> voice
    active = {}

    tick = 0
    tmp=0
    otimeusec=0

    for msg in track:
        tick += msg.time

        if msg.is_meta:
            if msg.type=="set_tempo":
                tmp=msg.tempo
                otimeusec+=msg.time*tmp
            continue
        else:
            otimeusec+=msg.time*tmp
        if msg.type=='sysex':#只有特定设备/软件能看明白的神秘信息，需要原样保留
            continue
        ch = msg.channel

        # 打击乐通道不处理
        if ch == 9:
            continue

        # -------- note_on --------
        if msg.type == 'note_on' and msg.velocity > 0:
            note = msg.note

            # 找一个真正“空闲”的 voice
            for v in voices[ch]:
                last = v[0]
                # end == 0 表示悬空音，不能复用
                if last.end != 0 and last.end <= tick:
                    break
            else:
                v = []
                voices[ch].append(v)

            ne = NoteEvent(
                start=tick,
                end=0,
                note=note,
                velocity=msg.velocity,
                orgchan=ch,
                timesec=otimeusec/tpb/1E6
            )
            v.append(ne)
            active[(ch, note)] = ne

        # -------- note_off / velocity=0 --------
        elif msg.type in ('note_off', 'note_on'):
            note = msg.note
            ne = active.pop((ch, note), None)
            if ne:
                ne.end = tick
            # 找不到就放过：脏 MIDI / 忽略
    print ("OK")
    return voices


def pack_voices(voices,max_channels_wish=6):
    print ("3. 通道重分配 ",end="")
    max_channels=88
    maxpolyphonycount=0
    maxpolyphonytick=[]
    """
    voices: split_channel_polyphony 的返回值
    max_channels: 最多可用的乐器通道数（不含打击乐）
    """
    # ---------- 1. 拉平所有 voice ----------
    all_voices = []
    for ch, vlist in voices.items():
        for v in vlist:
            if v:
                all_voices.append(v)

    # 按 voice 的起始时间排序（沙子从低处开始堆）
    all_voices.sort(key=lambda v: v[0].start)

    # ---------- 2. 全局压实 ----------
    events = []
    for voice in all_voices:
        for ne in voice:
            events.append(ne)

    events.sort(key=lambda ne: ne.start)

    packed_channels = [[] for _ in range(max_channels)]
    chan_end = [0] * max_channels
    packed_notes = []
    chanmap = defaultdict(set)
    chan_lastnote=[-1]* max_channels

    for ne in events:
        best = None
        addvoicemap=False
        for i in range(max_channels):
            if chan_end[i] <= ne.start:
                best = i
                for d in range(best,max_channels):
                    if chan_end[d] <= ne.start:
                        chan_lastnote[d]=-1
                maxpolyphonycount=max(maxpolyphonycount,i+1)
                if i+1>=max_channels_wish: 
                    addvoicemap=True
                    if maxpolyphonytick!=[]:
                        if maxpolyphonytick[-1][1]==ne.start:
                            maxpolyphonytick[-1]=[i+1,ne.start,ne.timesec,[]]
                        else:
                            maxpolyphonytick.append([i+1,ne.start,ne.timesec,[]])
                    else:
                        maxpolyphonytick.append([i+1,ne.start,ne.timesec,[]])
                break

        if best is None:
            #best=max_channels-1
            raise RuntimeError(
                "复音数超过max_channels限制\n"+
                "这是黑乐谱吗？？？\n"+
                "如果您设备性能强劲，可适当改大max_channels数值"
                )

        pn = PackedNote(
            start=ne.start,
            end=ne.end,
            note=ne.note,
            velocity=ne.velocity,
            outchan=best,
            orgchan=ne.orgchan
        )
        packed_notes.append(pn)
        chanmap[best].add(ne.orgchan)

        chan_end[best] = ne.end if ne.end != 0 else float("0")#对于某些原因找不到时长的音符的可选值：inf无限占用、0立即被占用、正整数占用固定时间
        
        chan_lastnote[best]=ne.note
        #print(chan_lastnote)
        #     C c D d E F f G g A a B
        plnt=[0,0,0,0,0,0,0,0,0,0,0,0]
        ntot=[[],[],[],[],[],[],[],[],[],[],[],[]]
        ntoto=[[],[],[],[],[],[],[],[],[],[],[],[]]
        for i in range(max_channels):
            if not chan_lastnote[i]==(-1):
                plnt[chan_lastnote[i]%12]+=1
                ntot[chan_lastnote[i]%12].append(chan_lastnote[i]//12)
        for i in range(12):
            for n in range(len(ntot[i])):
                if ntoto[i]==[]:
                    ntoto[i].append([ntot[i][n],1])
                else:
                    match=False
                    for q in range(len(ntoto[i])):
                        if ntoto[i][q][0]==ntot[i][n]:
                            ntoto[i][q][1]+=1
                            match=True
                    if not match:
                        ntoto[i].append([ntot[i][n],1])
            def ev_key(e):
                return (e[0])
            ntoto[i].sort(key=ev_key)
        if addvoicemap:
            maxpolyphonytick[-1][3]=ntoto
        #print(plnt,ntoto)
    # ---------- 3. 排序，准备输出 ----------
    packed_notes.sort(key=lambda n: (n.start, n.outchan))

    output = {
        "notes": packed_notes,
    }
    print ("OK")
    if maxpolyphonycount>max_channels_wish:
        print ("复音数超过max_channels_wish(",max_channels_wish,")，会出现缺音现象。\n请尝试在DAW中简化您的输入文件。")
        print ("\t 最大乐音复音数：",maxpolyphonycount)
        print ("\t 复音数\t Tick\t M:S.CS")
        for mpt in maxpolyphonytick:
            if mpt[0]>max_channels_wish:
                print ("\t",mpt[0],"\t",mpt[1],end="")
                st=mpt[2]
                msct=""
                msct+="{:d}".format(int(st//60))+":"
                msct+="{:0>2d}".format(int(st%60))+"."
                msct+="{:0>2d}".format(int((st%1)*100))
                print ("\t",msct)
                print (mpt[3])
    return output
    
def collect_orgchan_controls(track):
    print ("4. 抓取通道控制消息 ",end="")
    controls = defaultdict(list)
    tick = 0
    
    occ = OrgChanControl(tick=0,controls=[],aftertouch=[],pitchwheel=0,program=0)
    for msg in track:
        tick += msg.time
        if msg.is_meta or msg.type == 'sysex':
            continue
        ch = msg.channel
        if controls[ch]==[]:
            controls[ch].append(occ)
            ooc=copy.deepcopy(controls[ch][-1])
        else:
            ooc=copy.deepcopy(controls[ch][-1])
        if ch == 9:
            continue
        ooc.tick=tick
        if msg.type == 'program_change':
            ooc.program = msg.program

        elif msg.type == 'pitchwheel':
            ooc.pitchwheel = msg.pitch

        elif msg.type == 'control_change':
            if ooc.controls==[]:
                ooc.controls.append(msg)
            else:
                vleq=False
                for vl in ooc.controls:
                    if vl.control==msg.control:
                        vleq=True
                        vl.value=msg.value
                if vleq==False:
                    ooc.controls.append(msg)
        elif msg.type == 'aftertouch':
            ooc.aftertouch=msg
        else:
            ooc
        controls[ch].append(ooc)
    print ("OK")
    return controls

    
def embed_orgchan_controls(output, org_controls):
    print ("5. 控制消息重映射 ",end="")
    events = []
    def find_control_at(org_controls, tick,endtick):
        last = []
        for occ in org_controls:
            #惊天奇案
            if occ.tick < tick:
                continue
            else:
                last.append(occ)
                if occ.tick>=endtick:
                    break
                else:
                    continue
        return last
    # ---------- 1. 对每个音符匹配控制 ----------
    for n in output["notes"]:
        orgch = n.orgchan
        t = n.start
        t2= n.end
        
        occ = find_control_at(org_controls.get(orgch, []), t,t2)
        if occ:
            nocca=copy.deepcopy(occ)#谨防鬼故事
            for nocc in nocca:
                #nocc.tick=t#惊天奇案2
                events.append(("control", n.outchan, nocc))

        # 再加音符
        events.append(("note", t, n))

    # ---------- 2. 排序 ----------
    def ev_key(e):
        if e[0] == "control":
            return (e[2].tick, 0)   # 控制先于音符
        else:
            return (e[1], 1)

    events.sort(key=ev_key)
    print ("OK")
    return events


def expand_events(events):
    print ("6. 汇总控制消息 ",end="")
    out = []
    midilimitwarn=False
    for e in events:
        if e[0] == "note":
            _, tick, n = e
            ch=n.outchan
            if ch>=9:#谨防占用打击乐通道
                ch+=1
            else:
                ch
            if ch>15:
                if not midilimitwarn:
                    print ("复音数超过15，超过Midi通道限制，会出现缺音现象。\n请尝试在DAW中简化您的输入文件。")
                    midilimitwarn=True
                continue#丢弃处理——Midi最多15个乐音通道
            out.append((tick,
                mido.Message(
                    'note_on',
                    channel=ch,
                    note=n.note,
                    velocity=n.velocity,
                    time=0
                )
            ))

            if n.end > n.start:
                out.append((n.end,
                    mido.Message(
                        'note_off',
                        channel=ch,
                        note=n.note,
                        velocity=0,
                        time=0
                    )
                ))

        elif e[0] == "control":
            _, outch, occ = e
            t = occ.tick
            ch=outch
            if ch>=9:#谨防占用打击乐通道
                ch+=1
            else:
                ch
            if ch>15:
                if not midilimitwarn:
                    print ("\n复音数超过15，超过Midi通道限制，会出现缺音现象。\n请尝试在DAW中简化您的输入文件。")
                    midilimitwarn=True
                continue#丢弃处理
            if occ.program is not None:
                out.append((t,
                    mido.Message(
                        'program_change',
                        channel=ch,
                        program=occ.program,
                        time=0
                    )
                ))

            if occ.pitchwheel is not None:
                out.append((t,
                    mido.Message(
                        'pitchwheel',
                        channel=ch,
                        pitch=occ.pitchwheel,
                        time=0
                    )
                ))
            for msgs in occ.controls:
                if isinstance(msgs,list):
                    0
                else:
                    msgs=[msgs]
                for msg in msgs:
                    m = msg.copy(channel=ch, time=0)
                    out.append((t, m))
            if isinstance(occ.aftertouch,list):
                0
            else:
                m = occ.aftertouch.copy(channel=ch)
                out.append((t, m))
    print ("OK")
    return out

def merge_all_events(expanded,spev):
    all_events = []

    all_events.extend(expanded)
    all_events.extend(spev)
    all_events.sort(key=lambda x: x[0])
    return all_events

def collect_special_events(track):
    events = []
    tick = 0

    for msg in track:
        tick += msg.time
        if msg.is_meta:
            events.append((tick, msg))
            continue
        elif msg.type == "sysex":
            events.append((tick, msg))
            continue
        elif not msg.is_meta and msg.channel == 9:
            events.append((tick, msg))
            continue

    return events


def build_midi(events,spev, ticks_per_beat):
    print ("7. 构建Mido消息流 ",end="")
    mf = mido.MidiFile(ticks_per_beat=ticks_per_beat)
    track = mido.MidiTrack()
    mf.tracks.append(track)

    all_events = merge_all_events(events, spev)

    last_tick = 0
    for tick, msg in all_events:
        msg.time = tick - last_tick
        last_tick = tick
        track.append(msg)
    track.append(mido.MetaMessage('end_of_track', time=0))
    print ("OK")
    return mf


def compile_midi(path_in, path_out):
    mid = to_midi0(MidiFile(path_in))
    track = mid.tracks[0]

    voices = split_channel_polyphony(track,mid.ticks_per_beat)
    packed = pack_voices(voices)
    orgcontrols = collect_orgchan_controls(track)
    embed = embed_orgchan_controls(packed,orgcontrols)
    expan = expand_events(embed)
    spev = collect_special_events(track)
    out = build_midi(expan,spev,mid.ticks_per_beat)
    print ("8. 写出Midi文件 ",end="")
    out.save(path_out)
    print ("OK")

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python midisplitter2.py input.mid")
        sys.exit(1)
    
    if compile_midi(sys.argv[1],(sys.argv[1]+"_comp.mid")):
        sys.exit(0)
    else:
        sys.exit(1)